Shard load, full tensor sendaround

pytorch · Feb 6, 2025 · 65744ac · 65744ac
1 parent 4675681
commit 65744ac
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchdata/stateful_dataloader/ibm_rescalable.py b/torchdata/stateful_dataloader/ibm_rescalable.py
@@ -604,7 +604,7 @@ def load_distributed_state_dict(
     """
     base = loader.state_dict()
     nworkers = base["_snapshot"]["_main_snapshot"]["_num_workers"]
-    dstate = __pop_dstate(base, device_mesh, [dtensor.placement_types.Replicate()], True)
+    dstate = __pop_dstate(base, device_mesh, [dtensor.placement_types.Shard(0)], True)
     inp = {"state":deepcopy(base), "dstate":dstate}
     # Read distributed state dict
     reader = checkpoint.FileSystemReader(path)