From 4d6473384695f57621e8e7abe2ab220ba2c44579 Mon Sep 17 00:00:00 2001 From: Eugene Mironov Date: Tue, 11 Nov 2025 22:37:34 +0700 Subject: [PATCH] fixup! fixup! Add tests for flow matching models with RTC --- src/lerobot/policies/pi0/modeling_pi0.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/lerobot/policies/pi0/modeling_pi0.py b/src/lerobot/policies/pi0/modeling_pi0.py index c999e96c3..9b6f38ad4 100644 --- a/src/lerobot/policies/pi0/modeling_pi0.py +++ b/src/lerobot/policies/pi0/modeling_pi0.py @@ -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