Update on "Multiprocessing support for NT"

Fixes #110161 Allows NTs to be used in DataLoaders with `num_workers > 1`. [ghstack-poisoned]
pytorch · Oct 10, 2023 · a514685 · a514685
2 parents 3939e2f + da2fd0e
commit a514685
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 3 deletions.
diff --git a/tools/autograd/gen_python_functions.py b/tools/autograd/gen_python_functions.py
@@ -159,8 +159,6 @@
     "fill.Scalar",  # only used by the functionalization pass
     "lift.*",
     "normal_functional",  # only used by the functionalization pas
-    "_nested_view_from_buffer_copy",
-    "_nested_view_from_buffer_copy_out",
     "nbytes",
     "itemsize",
 ]

diff --git a/torch/multiprocessing/reductions.py b/torch/multiprocessing/reductions.py
@@ -358,7 +358,7 @@ def rebuild_nested_tensor(
     sizes = rebuild_sizes_func(*rebuild_sizes_args)
     strides = rebuild_strides_func(*rebuild_strides_args)
     offsets = rebuild_offsets_func(*rebuild_offsets_args)
-    return torch._nested_view_from_buffer(buffer, sizes, strides, offsets)
+    return torch._nested_view_from_buffer_copy(buffer, sizes, strides, offsets)
 
 
 def reduce_nested_tensor(nt):