Update torch.py

michaelfeil · Jan 30, 2024 · 065eab1 · 065eab1
1 parent 0ff0200
commit 065eab1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/libs/infinity_emb/infinity_emb/transformer/crossencoder/torch.py b/libs/infinity_emb/infinity_emb/transformer/crossencoder/torch.py
@@ -51,7 +51,7 @@ def __init__(self, model_name_or_path, **kwargs):
             self.model, logger, disable=self._target_device.type == "mps"
         )
 
-        if self.device.type == "cuda" and not os.environ.get(
+        if self._target_device.type == "cuda" and not os.environ.get(
             "INFINITY_DISABLE_HALF", ""
         ):
             logger.info(