Resolve MPS's lack of cummax

pytorch · Sep 4, 2024 · faa10a4 · faa10a4
1 parent 3658f95
commit faa10a4
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/ignite/metrics/vision/object_detection_average_precision_recall.py b/ignite/metrics/vision/object_detection_average_precision_recall.py
@@ -1,3 +1,4 @@
+import os
 from typing import Callable, cast, Dict, List, Optional, Sequence, Tuple, Union
 
 import torch
@@ -234,7 +235,9 @@ def _compute_average_precision(self, recall: torch.Tensor, precision: torch.Tens
         Returns:
             average_precision: (n-1)-dimensional tensor containing the average precision for mean dimensions.
         """
+        os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = True
         precision_integrand = precision.flip(-1).cummax(dim=-1).values.flip(-1)
+        os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = False
         rec_thresholds = cast(torch.Tensor, self.rec_thresholds).repeat((*recall.shape[:-1], 1))
         rec_thresh_indices = (
             torch.searchsorted(recall, rec_thresholds)