Update on "Pickle support for NT"

jbschlosser · jbschlosser · commit 2f9f7df7a16e · 2023-09-28T16:13:21.000-04:00
Fixes #104198 [ghstack-poisoned]
diff --git a/torch/_tensor.py b/torch/_tensor.py
@@ -376,7 +376,7 @@ def _reduce_ex_internal(self, proto):
                 self._nested_tensor_strides(),
                 self._nested_tensor_storage_offsets(),
             )
-            return (torch._utils._rebuild_nested_tensor, args_nested)
+            return (torch._nested_view_from_buffer, args_nested)
         elif (
             self.data_ptr() == 0
             and type(self) is not torch.Tensor
diff --git a/torch/_utils.py b/torch/_utils.py
@@ -304,12 +304,6 @@ def _rebuild_sparse_tensor(layout, data):
     raise NotImplementedError(f"rebuilding sparse tensor for layout {layout}")
 
 
-def _rebuild_nested_tensor(buffer, nested_sizes, nested_strides, storage_offsets):
-    return torch._nested_view_from_buffer(
-        buffer, nested_sizes, nested_strides, storage_offsets
-    )
-
-
 def _rebuild_device_tensor_from_numpy(data, dtype, device, requires_grad):
     tensor = torch.from_numpy(data).to(dtype=dtype, device=device)
     tensor.requires_grad = requires_grad
