fix stable diffusion rank error (#2208)

keras-team · Dec 1, 2023 · 2604105 · 2604105
1 parent d74a37a
commit 2604105
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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
                 )