[Bugfix] fix args with incontiguous node and edge tensor (#153)

python/dgl/backend/pytorch/tensor.py

@@ -137,7 +137,7 @@ def arange(start, stop):
     return th.arange(start, stop, dtype=th.int64)
 
 def zerocopy_to_dlpack(input):
-    return dlpack.to_dlpack(input)
+    return dlpack.to_dlpack(input.contiguous())
 
 def zerocopy_from_dlpack(dlpack_tensor):
     return dlpack.from_dlpack(dlpack_tensor)

tests/pytorch/test_index.py

@@ -28,8 +28,21 @@ def test_dlpack():
         x[0] = 1
         assert np.allclose(y.asnumpy(), ans)
 
+    def th2nd_incontiguous():
+        import dgl.backend as F
+
+        x = th.LongTensor([[0, 1], [2, 3]])
+        ans = np.array([0, 2])
+        y = x[:2, 0]
+        # Uncomment this line and comment the one below to observe error
+        #dl = dlpack.to_dlpack(y)
+        dl = F.zerocopy_to_dlpack(y)
+        z = nd.from_dlpack(dl)
+        assert np.allclose(z.asnumpy(), ans)
+
     nd2th()
     th2nd()
+    th2nd_incontiguous()
 
 def test_index():
     ans = np.ones((10,), dtype=np.int64) * 10