Merge branch 'main' into thomwolf_2024_06_17_readme_update

huggingface · Jun 18, 2024 · dba1e1a · dba1e1a
2 parents 4bcfa9f + b72d574
commit dba1e1a
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/lerobot/common/policies/diffusion/modeling_diffusion.py b/lerobot/common/policies/diffusion/modeling_diffusion.py
@@ -165,7 +165,9 @@ def __init__(self, config: DiffusionConfig):
         num_images = len([k for k in config.input_shapes if k.startswith("observation.image")])
         self.unet = DiffusionConditionalUnet1d(
             config,
-            global_cond_dim=(config.output_shapes["action"][0] + self.rgb_encoder.feature_dim * num_images)
+            global_cond_dim=(
+                config.input_shapes["observation.state"][0] + self.rgb_encoder.feature_dim * num_images
+            )
             * config.n_obs_steps,
         )