style

huggingface · Jan 10, 2025 · 032037b · 032037b
1 parent 1821dcf
commit 032037b
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/onnxruntime/test_diffusion.py b/tests/onnxruntime/test_diffusion.py
@@ -34,7 +34,7 @@
     ORTPipelineForInpainting,
     ORTPipelineForText2Image,
 )
-from optimum.utils import is_diffusers_version, is_transformers_version
+from optimum.utils import is_transformers_version
 from optimum.utils.testing_utils import grid_parameters, require_diffusers
 
 
@@ -222,8 +222,8 @@ def test_shape(self, model_arch: str):
                 self.assertEqual(outputs.shape, (batch_size, 3, height, width))
             else:
                 if model_arch == "flux":
-                    expected_height = height // 2**pipeline.vae_scale_factor
-                    expected_width = width // 2**pipeline.vae_scale_factor
+                    expected_height = height // (pipeline.vae_scale_factor * 2)
+                    expected_width = width // (pipeline.vae_scale_factor * 2)
                     channels = pipeline.transformer.config.in_channels
                     expected_shape = (batch_size, expected_height * expected_width, channels)
                 else: