fixed metrics device condition for mps (#693)

kaiko-ai · Oct 22, 2024 · 095ef7b · 095ef7b
1 parent b7fffbe
commit 095ef7b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/eva/core/models/modules/module.py b/src/eva/core/models/modules/module.py
@@ -53,7 +53,7 @@ def metrics_device(self) -> torch.device:
         device = os.getenv("METRICS_DEVICE", None)
         if device is not None:
             return torch.device(device)
-        elif self.device == torch.device("mps"):
+        elif self.device.type == "mps":
             # mps seems to have compatibility issues with segmentation metrics
             return torch.device("cpu")
         return self.device