fix

Lightning-AI · May 15, 2024 · 1170618 · 1170618
1 parent 82cbe27
commit 1170618
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/lightning/fabric/strategies/model_parallel.py b/src/lightning/fabric/strategies/model_parallel.py
@@ -484,7 +484,7 @@ def _load_checkpoint(
             raise ImportError("Loading a non-distributed checkpoint into a distributed model requires PyTorch >= 2.4.")
 
         checkpoint = torch.load(path, mmap=True, map_location="cpu")
-        _load_raw_module_state(checkpoint.pop(module_key), module, world_size=1, strict=strict)
+        _load_raw_module_state(checkpoint.pop(module_key), module, world_size=self.world_size, strict=strict)
 
         requested_metadata_keys = state.keys() - modules.keys() - optimizers.keys()
         _validate_keys_for_strict_loading(requested_metadata_keys, checkpoint.keys(), strict=strict)