diff options
| -rw-r--r-- | aten/src/ATen/DLConvertor.cpp | 7 | ||||
| -rw-r--r-- | aten/src/ATen/test/dlconvertor_test.cpp | 12 |
2 files changed, 19 insertions, 0 deletions
diff --git a/aten/src/ATen/DLConvertor.cpp b/aten/src/ATen/DLConvertor.cpp index e2983ebdff..a40e872724 100644 --- a/aten/src/ATen/DLConvertor.cpp +++ b/aten/src/ATen/DLConvertor.cpp @@ -178,6 +178,13 @@ Tensor fromDLPack(const DLManagedTensor* src) { auto deleter = [src](void* self) { src->deleter(const_cast<DLManagedTensor*>(src)); }; + if (!src->dl_tensor.strides) { + return at::from_blob(src->dl_tensor.data, + IntArrayRef(src->dl_tensor.shape, src->dl_tensor.ndim), + deleter, + at::device(device_type).dtype(stype)); + } + return at::from_blob( src->dl_tensor.data, IntArrayRef(src->dl_tensor.shape, src->dl_tensor.ndim), diff --git a/aten/src/ATen/test/dlconvertor_test.cpp b/aten/src/ATen/test/dlconvertor_test.cpp index e2c7a62403..4e5eb93df4 100644 --- a/aten/src/ATen/test/dlconvertor_test.cpp +++ b/aten/src/ATen/test/dlconvertor_test.cpp @@ -18,3 +18,15 @@ TEST(TestDlconvertor, TestDlconvertor) { ASSERT_TRUE(a.equal(b)); } + +TEST(TestDlconvertor, TestDlconvertorNoStrides) { + manual_seed(123); + + Tensor a = rand({3, 4}); + DLManagedTensor* dlMTensor = toDLPack(a); + dlMTensor->dl_tensor.strides = nullptr; + + Tensor b = fromDLPack(dlMTensor); + + ASSERT_TRUE(a.equal(b)); +} |