From 2604105858ac3e62b4ade89cb8d026f1c8feef62 Mon Sep 17 00:00:00 2001 From: Divyashree Sreepathihalli Date: Fri, 1 Dec 2023 12:05:21 -0800 Subject: [PATCH] fix stable diffusion rank error (#2208) --- keras_cv/models/stable_diffusion/stable_diffusion.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/keras_cv/models/stable_diffusion/stable_diffusion.py b/keras_cv/models/stable_diffusion/stable_diffusion.py index 975788ac74..768795c956 100644 --- a/keras_cv/models/stable_diffusion/stable_diffusion.py +++ b/keras_cv/models/stable_diffusion/stable_diffusion.py @@ -200,7 +200,7 @@ def generate_image( if diffusion_noise is not None: diffusion_noise = ops.squeeze(diffusion_noise) - if diffusion_noise.shape.rank == 3: + if len(ops.shape(diffusion_noise)) == 3: diffusion_noise = ops.repeat( ops.expand_dims(diffusion_noise, axis=0), batch_size, axis=0 )