switch to float32 for mac MPS

Deep-MI · Aug 25, 2024 · ac8128d · ac8128d
1 parent e2453c9
commit ac8128d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/HypVINN/inference.py b/HypVINN/inference.py
@@ -320,7 +320,7 @@ def eval(self, val_loader: DataLoader, pred_prob: torch.Tensor, out_scale: float
 
             images = batch["image"].to(self.device)
             scale_factors = batch["scale_factor"].to(self.device)
-            weight_factors = batch["weight_factor"].to(self.device)
+            weight_factors = batch["weight_factor"].to(self.device, dtype=torch.float32)
 
             pred = self.model(images, scale_factors, weight_factors, out_scale)