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

refactor pi052 to reuse pi05

Browse Source
This commit is contained in:
Pepijn
2026-07-15 19:26:55 +02:00
parent 55d9ff740e
commit 5b8e6ffe8e
5 changed files with 223 additions and 563 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/policies/pi052/test_pi052_vision_checkpointing.py
+6 -8
View File
@@ -41,14 +41,12 @@ def _checkpoint_model():
tower = _MockVisionTower()
language_model = SimpleNamespace(gradient_checkpointing=False)
expert_model = SimpleNamespace(gradient_checkpointing=False)
model = SimpleNamespace(
gradient_checkpointing_enabled=False,
paligemma_with_expert=SimpleNamespace(
paligemma=SimpleNamespace(
model=SimpleNamespace(language_model=language_model, vision_tower=tower)
),
gemma_expert=SimpleNamespace(model=expert_model),
),
model = PI05Pytorch.__new__(PI05Pytorch)
nn.Module.__init__(model)
model.gradient_checkpointing_enabled = False
model.paligemma_with_expert = SimpleNamespace(
paligemma=SimpleNamespace(model=SimpleNamespace(language_model=language_model, vision_tower=tower)),
gemma_expert=SimpleNamespace(model=expert_model),
)
return model, tower, language_model, expert_model