Handle crop_shape=None in Diffusion Policy (#219)

2024-05-28 18:27:33 +01:00
parent e3b9f1c19b
commit 3d625ae6d3
3 changed files with 8 additions and 6 deletions
--- a/lerobot/common/policies/diffusion/configuration_diffusion.py
+++ b/lerobot/common/policies/diffusion/configuration_diffusion.py
@@ -155,7 +155,7 @@ class DiffusionConfig:
                f"{self.__class__.__name__} only handles one image for now. Got image keys {image_keys}."
            )
        image_key = next(iter(image_keys))
-        if (
+        if self.crop_shape is not None and (
            self.crop_shape[0] > self.input_shapes[image_key][1]
            or self.crop_shape[1] > self.input_shapes[image_key][2]
        ):