avoid loading the model to rank 0 to avoid a big vram spike which can OOM

src/lerobot/policies/pretrained.py

@@ -29,6 +29,7 @@ from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download, sa
 from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
 from huggingface_hub.errors import HfHubHTTPError
 from safetensors.torch import load_model as load_model_as_safetensor, save_model as save_model_as_safetensor
+import torch
 from torch import Tensor, nn
 
 from lerobot.__version__ import __version__
@@ -221,6 +222,14 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
 
     @classmethod
     def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
+        # safetensors' load_file maps the bare string "cuda" to cuda:0 regardless of the current
+        # device (unlike torch's .to("cuda"), which honors torch.cuda.current_device()). Under
+        # multi-GPU accelerate/FSDP every rank would then load its weights onto GPU 0, OOMing it
+        # before sharding. Resolve "cuda" to the concrete current-device index so each rank loads
+        # onto its own GPU.
+        if map_location == "cuda" and torch.cuda.is_available():
+            map_location = f"cuda:{torch.cuda.current_device()}"
+
         # Create base kwargs
         kwargs = {"strict": strict}