Merge branch 'claude/fix-blackwell-monai-tests-Apjwl' into 8587-test-erros-on-pytorch-release-2508-on-series-50

Author: garciadias

monai/networks/blocks/warp.py

@@ -64,6 +64,18 @@ def __init__(self, mode=GridSampleMode.BILINEAR.value, padding_mode=GridSamplePa
         super().__init__()
         # resolves _interp_mode for different methods
 
+        # Native string modes are always stored for the PyTorch fallback path.
+        # When USE_COMPILED=True but the device is unsupported at runtime (e.g. Blackwell GPU),
+        # forward() falls back to F.grid_sample which requires string, not integer, modes.
+        self._interp_mode_native = (
+            GridSampleMode(mode).value if mode in (m.value for m in GridSampleMode) else GridSampleMode.BILINEAR.value
+        )
+        self._padding_mode_native = (
+            GridSamplePadMode(padding_mode).value
+            if padding_mode in (p.value for p in GridSamplePadMode)
+            else GridSamplePadMode.BORDER.value
+        )
+
         if USE_COMPILED:
             if mode in (inter.value for inter in GridSampleMode):
                 mode = GridSampleMode(mode)
@@ -78,7 +90,7 @@ def __init__(self, mode=GridSampleMode.BILINEAR.value, padding_mode=GridSamplePa
             self._interp_mode = mode
         else:
             warnings.warn("monai.networks.blocks.Warp: Using PyTorch native grid_sample.")
-            self._interp_mode = GridSampleMode(mode).value
+            self._interp_mode = self._interp_mode_native
 
         # resolves _padding_mode for different methods
         if USE_COMPILED:
@@ -94,7 +106,7 @@ def __init__(self, mode=GridSampleMode.BILINEAR.value, padding_mode=GridSamplePa
                     padding_mode = 0  # default to nearest
             self._padding_mode = padding_mode
         else:
-            self._padding_mode = GridSamplePadMode(padding_mode).value
+            self._padding_mode = self._padding_mode_native
 
         self.ref_grid = None
         self.jitter = jitter
@@ -147,7 +159,7 @@ def forward(self, image: torch.Tensor, ddf: torch.Tensor):
             index_ordering: list[int] = list(range(spatial_dims - 1, -1, -1))
             grid = grid[..., index_ordering]  # z, y, x -> x, y, z
             return F.grid_sample(
-                image, grid, mode=self._interp_mode, padding_mode=f"{self._padding_mode}", align_corners=True
+                image, grid, mode=self._interp_mode_native, padding_mode=self._padding_mode_native, align_corners=True
             )
 
         # using csrc resampling

tests/transforms/test_spatial_gpu_support.py

@@ -14,6 +14,7 @@
 from __future__ import annotations
 
 import unittest
+from unittest.mock import MagicMock, patch
 
 import torch
 
@@ -28,23 +29,16 @@ def test_cpu_device_always_supported(self):
         device = torch.device("cpu")
         self.assertFalse(_compiled_unsupported(device))
 
-    def test_non_cuda_device_always_supported(self):
-        """Non-CUDA devices should always be supported."""
-        device = torch.device("cpu")
-        self.assertFalse(_compiled_unsupported(device))
-
     @unittest.skipIf(not torch.cuda.is_available(), reason="CUDA not available")
     def test_cuda_device_detection(self):
         """Verify CUDA compute capability detection."""
-        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-        if device.type == "cuda":
-            cc_major = torch.cuda.get_device_properties(device).major
-            unsupported = _compiled_unsupported(device)
-            # Device is unsupported if cc_major >= 12
-            if cc_major >= 12:
-                self.assertTrue(unsupported)
-            else:
-                self.assertFalse(unsupported)
+        device = torch.device("cuda:0")
+        cc_major = torch.cuda.get_device_properties(device).major
+        unsupported = _compiled_unsupported(device)
+        if cc_major >= 12:
+            self.assertTrue(unsupported)
+        else:
+            self.assertFalse(unsupported)
 
     def test_compiled_unsupported_return_type(self):
         """Verify return type is bool."""
@@ -56,19 +50,24 @@ def test_compiled_unsupported_return_type(self):
 class TestResampleFallback(unittest.TestCase):
     """Test Resample fallback behavior on unsupported devices."""
 
-    @unittest.skipIf(not torch.cuda.is_available(), reason="CUDA not available")
     def test_resample_compilation_flag_respected(self):
-        """Verify Resample respects _compiled_unsupported check."""
-        # This would require internal inspection or output verification
-        # Could test with mock device properties or actual Blackwell GPU
+        """Verify _compiled_unsupported identifies Blackwell (cc>=12) and supported (cc<12) devices."""
+        mock_props = MagicMock()
+        cuda_device = torch.device("cuda:0")
+
+        mock_props.major = 12  # Blackwell – unsupported
+        with patch("torch.cuda.get_device_properties", return_value=mock_props):
+            self.assertTrue(_compiled_unsupported(cuda_device))
+
+        mock_props.major = 9  # Hopper – supported
+        with patch("torch.cuda.get_device_properties", return_value=mock_props):
+            self.assertFalse(_compiled_unsupported(cuda_device))
 
     def test_compiled_unsupported_logic(self):
         """Test that unsupported devices are correctly detected."""
-        # CPU should be supported
         cpu_device = torch.device("cpu")
         self.assertFalse(_compiled_unsupported(cpu_device))
 
-        # Verify logic: return True if CUDA and cc_major >= 12
         cuda_device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
         if cuda_device.type == "cuda":
             cc_major = torch.cuda.get_device_properties(cuda_device).major
@@ -79,5 +78,3 @@ def test_compiled_unsupported_logic(self):
 
 if __name__ == "__main__":
     unittest.main()
-if __name__ == "__main__":
-    unittest.main()