fix

Author: awaelchli

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)