add mps support to utils.to_numpy

skorch-dev · Aug 8, 2022 · 9847ea4 · 9847ea4
1 parent 64a2d16
commit 9847ea4
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/skorch/tests/test_utils.py b/skorch/tests/test_utils.py
@@ -199,6 +199,16 @@ def test_invalid_inputs(self, to_numpy, x_invalid):
         expected = "Cannot convert this data type to a numpy array."
         assert e.value.args[0] == expected
 
+    @pytest.mark.skipif(
+        not (hasattr(torch.backends, "mps") and torch.backends.mps.is_available()),
+        reason='Skipped because mps is not available as a torch backend'
+    )
+    def test_mps_support(self, to_numpy, x_tensor):
+        device = torch.device('mps')
+        x_tensor.to(device)
+        x_numpy = to_numpy(x_tensor)
+        self.compare_array_to_tensor(x_numpy, x_tensor)
+
 
 class TestToDevice:
     @pytest.fixture

diff --git a/skorch/utils.py b/skorch/utils.py
@@ -165,6 +165,9 @@ def to_numpy(X):
     if X.is_cuda:
         X = X.cpu()
 
+    if hasattr(X, 'is_mps') and X.is_mps:
+        X = X.cpu()
+
     if X.requires_grad:
         X = X.detach()