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

fixup! fixup! Add tests for flow matching models with RTC

Browse Source
This commit is contained in:
Eugene Mironov
2025-11-11 22:37:34 +07:00
parent 0c3ed6ca7a
commit 4d64733846
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/policies/pi0/modeling_pi0.py
+4 -4
View File
@@ -916,8 +916,8 @@ class PI0Policy(PreTrainedPolicy):
self.config = config
# Initialize the core PI0 model
self.model = PI0Pytorch(config, rtc_processor=self.rtc_processor)
self.init_rtc_processor()
self.model = PI0Pytorch(config, rtc_processor=self.rtc_processor)
# Enable gradient checkpointing if requested
if config.gradient_checkpointing:
@@ -1116,9 +1116,9 @@ class PI0Policy(PreTrainedPolicy):
if self.config.rtc_config is not None:
self.rtc_processor = RTCProcessor(self.config.rtc_config)
# Set rtc_processor to the model if it exists
if self.model is not None:
self.model.rtc_processor = self.rtc_processor
model_value = getattr(self, "model", None)
if model_value is not None:
model_value.rtc_processor = self.rtc_processor
def _rtc_enabled(self) -> bool:
return self.config.rtc_config is not None and self.config.rtc_config.enabled