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`