Relax stride check in unit dimensions.

dali/python/backend_impl.cc

@@ -181,19 +181,18 @@ void CheckContiguousTensor(const TStrides &strides, int num_strides,
                            const TShape &shape, int num_extents, size_t element_size) {
   DALI_ENFORCE(num_strides == num_extents,
     "There should be exactly as many strides as there are extents in array shape.");
+  for (int i = 0; i < num_extents; i++)
+    if (shape[i] == 0)
+      return;  // The volume is 0, the strides will never be used to compute an actual address
+
   int64_t stride_from_shape = element_size;
-  int64_t stride_from_shape_collapsed = 1;
   int64_t last_non_one_dim = 1;
   for (int i = num_strides - 1; i >= 0; i--) {
-    DALI_ENFORCE(strides[i] == stride_from_shape || strides[i] == stride_from_shape_collapsed,
+    DALI_ENFORCE(shape[i] == 1 ||  // ignore unit extents - the stride won't be used anyway
+                 strides[i] == stride_from_shape,
         make_string("Strided data not supported. Dimension ", i, " has stride ", strides[i],
         " whereas densely packed data of this shape would have a stride ", stride_from_shape));
     stride_from_shape *= shape[i];
-    // for shapes [1, 1, 5] leading dimensions may not contribute to stride
-    if (shape[i] != 1) {
-      stride_from_shape_collapsed *= last_non_one_dim;
-      last_non_one_dim = shape[i];
-    }
   }
 }
 

dali/test/python/experimental_mode/test_interop.py

@@ -89,6 +89,20 @@ def test_torch_nocopy(device: str):
     np.testing.assert_array_equal(ndd_tensor.cpu(), [42, 2, 3])
 
 
+@attr("pytorch")
+@params(("cpu",), ("gpu",))
+def test_from_torch(device: str):
+    import torch
+
+    if version.parse(torch.__version__) < version.parse("2.6.0"):
+        raise SkipTest("Requires PyTorch >= 2.6.0")
+
+    inpt = torch.arange(1, 5).reshape((1, 2, 2))
+    ndd_tensor = ndd.as_tensor(inpt)
+
+    np.testing.assert_array_equal(ndd_tensor.cpu(), [[[1, 2], [3, 4]]])
+
+
 @attr("pytorch")
 @params(("cpu",), ("gpu",))
 def test_tensor_to_torch(device: str):