From c5a029a28a7e6aa701d199da1f608b8ee32dba25 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Sat, 13 Sep 2025 11:12:54 +0200 Subject: [PATCH] also compile forward method --- src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py b/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py index 120791cc1..549dc0a9b 100644 --- a/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py +++ b/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py @@ -517,6 +517,8 @@ class PI0Pytorch(nn.Module): # see openpi `PI0Pytorch` if config.compile_model: torch.set_float32_matmul_precision("high") self.sample_actions = torch.compile(self.sample_actions, mode=config.compile_mode) + # Also compile the main forward pass used during training + self.forward = torch.compile(self.forward, mode=config.compile_mode) msg = """transformers_replace is not installed correctly. Please install it with `pip install transformers==4.53.2`