diff --git a/src/lerobot/policies/sarm/modeling_sarm.py b/src/lerobot/policies/sarm/modeling_sarm.py index d0a49fdd5..4cd57dced 100644 --- a/src/lerobot/policies/sarm/modeling_sarm.py +++ b/src/lerobot/policies/sarm/modeling_sarm.py @@ -309,7 +309,7 @@ class SARMRewardModel(PreTrainedPolicy): self.config.num_stages = num_stages self.config.subtask_names = subtask_names - logging.info(f"Auto-detected {num_stages} subtasks from dataset: {subtask_names}") + logging.info(f"Auto-detected {num_stages} subtasks from dataset: {subtask_names}, using {num_stages} stages") def to(self, device): """Override to method to ensure all components move together."""