fix the torch dtype check

2024-04-11 08:56:46 -04:00
parent 5767eea874
commit 4c92b51cd5
1 changed files with 2 additions and 2 deletions
--- a/tests/e2e/test_lora_llama.py
+++ b/tests/e2e/test_lora_llama.py
@@ -125,6 +125,6 @@ class TestLoraLlama(unittest.TestCase):
            config = f_handle.read()
        config = json.loads(config)
        if is_torch_bf16_gpu_available():
-            assert config["torch_dtype"] == "torch.bfloat16"
+            assert config["torch_dtype"] == "bfloat16"
        else:
-            assert config["torch_dtype"] == "torch.float16"
+            assert config["torch_dtype"] == "float16"