admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-24 10:16:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

also compile forward method

Browse Source
This commit is contained in:
Pepijn
2025-09-13 11:12:54 +02:00
parent c8163662ad
commit c5a029a28a
1 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py
+2
View File
@@ -517,6 +517,8 @@ class PI0Pytorch(nn.Module): # see openpi `PI0Pytorch`
if config.compile_model: if config.compile_model:
torch.set_float32_matmul_precision("high") torch.set_float32_matmul_precision("high")
self.sample_actions = torch.compile(self.sample_actions, mode=config.compile_mode) 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. msg = """transformers_replace is not installed correctly.
Please install it with `pip install transformers==4.53.2` Please install it with `pip install transformers==4.53.2`