Nested fix

pytorch · Jun 16, 2022 · c534537 · c534537
1 parent f965681
commit c534537
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 2 deletions.
diff --git a/aten/src/ATen/core/tensor_type.cpp b/aten/src/ATen/core/tensor_type.cpp
@@ -253,7 +253,7 @@ TensorTypePtr TensorType::create(const at::Tensor& t) {
   VaryingShape<size_t> stride_indices;
   VaryingShape<int64_t> strides;
   VaryingShape<int64_t> sizes;
-  if (t.layout() == at::kStrided) {
+  if (t.layout() == at::kStrided && !t.is_nested()) {
     sizes = VaryingShape<int64_t>{t.sizes().vec()};
     strides = VaryingShape<int64_t>{t.strides().vec()};
     return TensorType::create(

diff --git a/torch/csrc/jit/ir/constants.cpp b/torch/csrc/jit/ir/constants.cpp
@@ -12,7 +12,7 @@ namespace jit {
 bool insertableTensor(const at::Tensor& ten) {
   // bail if tensor has no storage i.e. opaque tensor used in MKLdnn.
   // or gradients because we have no way of serializing them & are mutable
-  return !ten.requires_grad() && ten.has_storage();
+  return !ten.requires_grad() && ten.has_storage() && !ten.is_nested();
 }
 
 bool insertableIValue(const IValue& ivalue) {

diff --git a/torch/csrc/jit/ir/node_hashing.cpp b/torch/csrc/jit/ir/node_hashing.cpp
@@ -23,6 +23,9 @@ bool tensorEqual(const at::Tensor& lhs, const at::Tensor& rhs) {
   if (lhs.is_mkldnn() || rhs.is_mkldnn()) {
     return false;
   }
+  if (lhs.is_nested() || rhs.is_nested()) {
+    return false;
+  }
   // If device is not equal, lhs.equal(rhs) would throw an error.
   if (lhs.device() != rhs.device()) {
     return false;