add for rest of policies

src/lerobot/policies/pi0/modeling_pi0.py

@@ -995,6 +995,13 @@ class PI0Policy(PreTrainedPolicy):
 
         # Initialize model without loading weights
         # Check if dataset_stats were provided in kwargs
+        if _transformers_available:
+            from transformers.modeling_utils import no_init_weights
+
+            with no_init_weights():
+                model = cls(config, **kwargs)
+            model.model.paligemma_with_expert.paligemma.tie_weights()
+        else:
             model = cls(config, **kwargs)
 
         # Now manually load and remap the state dict

src/lerobot/policies/pi05/modeling_pi05.py

@@ -969,6 +969,7 @@ class PI05Policy(PreTrainedPolicy):
         # Check if dataset_stats were provided in kwargs
         if _transformers_available:
             from transformers.modeling_utils import no_init_weights
+
             with no_init_weights():
                 model = cls(config, **kwargs)
             model.model.paligemma_with_expert.paligemma.tie_weights()

src/lerobot/policies/pi0_fast/modeling_pi0_fast.py

@@ -895,6 +895,13 @@ class PI0FastPolicy(PreTrainedPolicy):
 
         # Initialize model without loading weights
         # Check if dataset_stats were provided in kwargs
+        if _transformers_available:
+            from transformers.modeling_utils import no_init_weights
+
+            with no_init_weights():
+                model = cls(config, **kwargs)
+            model.model.paligemma_with_expert.paligemma.tie_weights()
+        else:
             model = cls(config, **kwargs)
 
         # Now manually load and remap the state dict