From 2af40615b8010ed56b8d6a71b4d9221c22db4cb8 Mon Sep 17 00:00:00 2001
From: Pepijn <pepijn@huggingface.co>
Date: Tue, 25 Nov 2025 14:48:52 +0100
Subject: [PATCH] add image validation

---
 src/lerobot/policies/sarm/configuration_sarm.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/src/lerobot/policies/sarm/configuration_sarm.py b/src/lerobot/policies/sarm/configuration_sarm.py
index 437d4bb59..624c5ffd0 100644
--- a/src/lerobot/policies/sarm/configuration_sarm.py
+++ b/src/lerobot/policies/sarm/configuration_sarm.py
@@ -96,6 +96,13 @@ class SARMConfig(PreTrainedConfig):
     def __post_init__(self):
         super().__post_init__()
         
+        # Add the image_key, the processor will transform this into video_features
+        if self.image_key and self.image_key not in self.input_features:
+            self.input_features[self.image_key] = PolicyFeature(
+                shape=(480, 640, 3),
+                type=FeatureType.VISUAL
+            )
+        
         # Validate configuration
         if self.hidden_dim % self.num_heads != 0:
             raise ValueError(