admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-05-22 03:59:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix pi05 forward compile (#2551)

Browse Source
This commit is contained in:
Michel Aractingi
2025-12-02 11:01:43 +01:00
committed by GitHub
parent af4766b602
commit 797cd2725a
1 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/policies/pi05/modeling_pi05.py
+2
View File
@@ -538,6 +538,8 @@ class PI05Pytorch(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 = """An incorrect transformer version is used, please create an issue on https://github.com/huggingface/lerobot/issues""" msg = """An incorrect transformer version is used, please create an issue on https://github.com/huggingface/lerobot/issues"""