[Bugfix] Allow pinning of empty tensor (#6044)

dmlc · Jul 25, 2023 · c08192c · c08192c
1 parent b400af4
commit c08192c
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 5 deletions.
diff --git a/python/dgl/backend/pytorch/tensor.py b/python/dgl/backend/pytorch/tensor.py
@@ -451,11 +451,13 @@ def check_is_view(input):
 
 
 def zerocopy_to_dgl_ndarray_for_write(input):
-    assert input.is_contiguous(), (
-        "Cannot convert non-contiguous tensors "
-        "to dgl ndarray for write. Call .to_contiguous() first."
-    )
-    check_is_view(input)
+    if input.numel() > 0:
+        # only check non-empty tensors
+        assert input.is_contiguous(), (
+            "Cannot convert non-contiguous tensors "
+            "to dgl ndarray for write. Call .to_contiguous() first."
+        )
+        check_is_view(input)
     return zerocopy_to_dgl_ndarray(input)
 
 

diff --git a/tests/python/pytorch/utils/test_pin_memory.py b/tests/python/pytorch/utils/test_pin_memory.py
@@ -31,6 +31,10 @@ def test_pin_view():
     with pytest.raises(dgl.DGLError):
         dgl.utils.pin_memory_inplace(v)
 
+    # make sure an empty view does not generate an error
+    u = t[10:10]
+    u = dgl.utils.pin_memory_inplace(u)
+
 
 @pytest.mark.skipif(
     F._default_context_str == "cpu", reason="Need gpu for this test."