fixup! Add tests for flow matching models with RTC

This commit is contained in:
Eugene Mironov
2025-11-11 22:30:57 +07:00
parent a59ebab66b
commit a14e8a65cd

View File

@@ -916,8 +916,8 @@ class PI0Policy(PreTrainedPolicy):
self.config = config
# Initialize the core PI0 model
self.init_rtc_processor()
self.model = PI0Pytorch(config, rtc_processor=self.rtc_processor)
self.init_rtc_processor()
# Enable gradient checkpointing if requested
if config.gradient_checkpointing: