diff --git a/optimum/intel/openvino/modeling_diffusion.py b/optimum/intel/openvino/modeling_diffusion.py
index e86c5a8f0..8cfc9529a 100644
--- 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: