Use float32 in metrics when metric device is MPS

Related to #3326

Author: vfdev-5

ignite/metrics/multilabel_confusion_matrix.py

@@ -136,7 +136,11 @@ def compute(self) -> torch.Tensor:
             raise NotComputableError("Confusion matrix must have at least one example before it can be computed.")
 
         if self.normalized:
-            conf = self.confusion_matrix.to(dtype=torch.float64)
+            # MPS framework doesn't support float64, should use float32
+            double_dtype = torch.float64
+            if self.confusion_matrix.device.type == "mps":
+                double_dtype = torch.float32
+            conf = self.confusion_matrix.to(dtype=double_dtype)
             sums = conf.sum(dim=(1, 2))
             return conf / sums[:, None, None]
 

ignite/metrics/precision.py

@@ -38,6 +38,10 @@ def __init__(
         super(_BasePrecisionRecall, self).__init__(
             output_transform=output_transform, is_multilabel=is_multilabel, device=device, skip_unrolling=skip_unrolling
         )
+        # MPS framework doesn't support float64, should use float32
+        self._double_dtype = torch.float64
+        if self._device.type == "mps":
+            self._double_dtype = torch.float32
 
     def _check_type(self, output: Sequence[torch.Tensor]) -> None:
         super()._check_type(output)
@@ -81,8 +85,8 @@ def _prepare_output(self, output: Sequence[torch.Tensor]) -> Sequence[torch.Tens
             y = torch.transpose(y, 1, -1).reshape(-1, num_labels)
 
         # Convert from int cuda/cpu to double on self._device
-        y_pred = y_pred.to(dtype=torch.float64, device=self._device)
-        y = y.to(dtype=torch.float64, device=self._device)
+        y_pred = y_pred.to(dtype=self._double_dtype, device=self._device)
+        y = y.to(dtype=self._double_dtype, device=self._device)
         correct = y * y_pred
 
         return y_pred, y, correct