Fix reshaping unet if timestep is 0d tensor (#1083)

huggingface · Dec 19, 2024 · cda4908 · cda4908
1 parent 106a5b7
commit cda4908
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/optimum/intel/openvino/modeling_diffusion.py b/optimum/intel/openvino/modeling_diffusion.py
@@ -657,7 +657,8 @@ def _reshape_unet(
         for inputs in model.inputs:
             shapes[inputs] = inputs.get_partial_shape()
             if inputs.get_any_name() == "timestep":
-                shapes[inputs][0] = 1
+                if shapes[inputs].rank == 1:
+                    shapes[inputs][0] = 1
             elif inputs.get_any_name() == "sample":
                 in_channels = self.unet.config.get("in_channels", None)
                 if in_channels is None: