Get test passing on devbox with GPU. Add additional test for requires…

…_grad AND gpu
anthropics · Jan 16, 2025 · 6ca9740 · 6ca9740
1 parent c399f6a
commit 6ca9740
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 6 deletions.
diff --git a/src/serialize/per_type/pytorch.rs b/src/serialize/per_type/pytorch.rs
@@ -36,16 +36,26 @@ impl<'a> Serialize for PyTorchSerializer<'a> {
             let detached = PyObject_CallMethodObjArgs(self.previous.ptr, detach_method, std::ptr::null_mut::<pyo3_ffi::PyObject>());
             Py_DECREF(detach_method);
 
-            // Get numpy() method from detached tensor
-            let numpy_method = PyUnicode_InternFromString("numpy\0".as_ptr() as *const c_char);
-            let numpy_array = if detached.is_null() {
-                // If detach failed (tensor doesn't require grad), try numpy directly
-                PyObject_CallMethodObjArgs(self.previous.ptr, numpy_method, std::ptr::null_mut::<pyo3_ffi::PyObject>())
+            // Get cpu() method to ensure tensor is on CPU
+            let cpu_method = PyUnicode_InternFromString("cpu\0".as_ptr() as *const c_char);
+            let cpu_tensor = if detached.is_null() {
+                PyObject_CallMethodObjArgs(self.previous.ptr, cpu_method, std::ptr::null_mut::<pyo3_ffi::PyObject>())
             } else {
-                let result = PyObject_CallMethodObjArgs(detached, numpy_method, std::ptr::null_mut::<pyo3_ffi::PyObject>());
+                let result = PyObject_CallMethodObjArgs(detached, cpu_method, std::ptr::null_mut::<pyo3_ffi::PyObject>());
                 Py_DECREF(detached);
                 result
             };
+            Py_DECREF(cpu_method);
+
+            // Get numpy() method from CPU tensor
+            let numpy_method = PyUnicode_InternFromString("numpy\0".as_ptr() as *const c_char);
+            let numpy_array = if !cpu_tensor.is_null() {
+                let result = PyObject_CallMethodObjArgs(cpu_tensor, numpy_method, std::ptr::null_mut::<pyo3_ffi::PyObject>());
+                Py_DECREF(cpu_tensor);
+                result
+            } else {
+                std::ptr::null_mut()
+            };
             Py_DECREF(numpy_method);
 
             if numpy_array.is_null() {

diff --git a/test/test_pytorch.py b/test/test_pytorch.py
@@ -71,6 +71,15 @@ def test_tensor_on_gpu(self):
         tensor = torch.tensor([1, 2, 3]).cuda()
         self.assertEqual(orjson.dumps(tensor, option=orjson.OPT_SERIALIZE_NUMPY), b'[1,2,3]')
 
+    def test_tensor_on_gpu_and_requires_grad(self):
+        """
+        torch.Tensor on GPU if available AND requires_grad=True
+        """
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA not available")
+        tensor = torch.tensor([1., 2., 3.], requires_grad=True).cuda()
+        self.assertEqual(orjson.dumps(tensor, option=orjson.OPT_SERIALIZE_NUMPY), b'[1.0,2.0,3.0]')
+
     def test_tensor_zero_dim(self):
         """
         Test 0-dimensional tensors are properly serialized as scalar values