Solve the qwix bugs when config.per_device_batch_size * jax.local_device_count() < 1 (#237)

* fix wan unit test bugs

* line problems

* bug fix under low per_device_batch_size

Author: susanbao

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -286,7 +286,7 @@ def quantize_transformer(cls, config: HyperParameters, model: WanModel, pipeline
       return model
     max_logging.log("Quantizing transformer with Qwix.")
 
-    batch_size = int(config.per_device_batch_size * jax.local_device_count())
+    batch_size = jnp.ceil(config.per_device_batch_size * jax.local_device_count()).astype(jnp.int32)
     latents, prompt_embeds, timesteps = get_dummy_wan_inputs(config, pipeline, batch_size)
     model_inputs = (latents, timesteps, prompt_embeds)
     with mesh: