[BugFix] Remove forkingpickler (#1049)

pytorch · Oct 21, 2024 · 1659518 · 1659518
1 parent b2ccfe7
commit 1659518
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tensordict/_reductions.py b/tensordict/_reductions.py
@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 import copyreg
-from multiprocessing.reduction import ForkingPickler
+from multiprocessing import reduction
 
 import torch
 from tensordict._lazy import LazyStackedTensorDict
@@ -160,10 +160,10 @@ def _reduce_td(data: TensorDict):
     # return (_rebuild_tensordict_files, (flat_key_values, metadata_dict))
 
 
-ForkingPickler.register(TensorDict, _reduce_td)
+reduction.register(TensorDict, _reduce_td)
 
 copyreg.pickle(TensorDict, _reduce_td)
 
-ForkingPickler.register(LazyStackedTensorDict, _reduce_td)
+reduction.register(LazyStackedTensorDict, _reduce_td)
 
 copyreg.pickle(LazyStackedTensorDict, _reduce_td)