From 49444652c66a2da98f9b122d0fbb5f9c8e3c20d5 Mon Sep 17 00:00:00 2001 From: Jeremiah Coholich Date: Fri, 20 Feb 2026 15:41:44 -0500 Subject: [PATCH] speedup pi-05 modeling loading by 72s --- src/lerobot/policies/pi05/modeling_pi05.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/lerobot/policies/pi05/modeling_pi05.py b/src/lerobot/policies/pi05/modeling_pi05.py index 104ec63bf..3929e9ebf 100644 --- a/src/lerobot/policies/pi05/modeling_pi05.py +++ b/src/lerobot/policies/pi05/modeling_pi05.py @@ -967,7 +967,13 @@ class PI05Policy(PreTrainedPolicy): # Initialize model without loading weights # Check if dataset_stats were provided in kwargs - model = cls(config, **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 try: