From 8ef3f08cde81c5523b0f5f636ed5e0be732f41fa Mon Sep 17 00:00:00 2001 From: eaidova Date: Thu, 19 Dec 2024 15:36:02 +0400 Subject: [PATCH] fix reshaping unet if timestep is 0d tensor --- optimum/intel/openvino/modeling_diffusion.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/optimum/intel/openvino/modeling_diffusion.py b/optimum/intel/openvino/modeling_diffusion.py index e86c5a8f02..8cfc9529a2 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: