Update app.py

app.py

@@ -122,7 +122,8 @@ if len(sample_batch['pixel_values']) > 0:
 def train_step(state, batch, rng):
     def compute_loss(params, pixel_values, rng):
         pixel_values = jnp.array(pixel_values, dtype=jnp.float32)
-        pixel_values = jnp.expand_dims(pixel_values, axis=0)  # Add batch dimension
+        if pixel_values.ndim == 3:
+            pixel_values = jnp.expand_dims(pixel_values, axis=0)  # Add batch dimension if needed
         print(f"pixel_values shape in compute_loss: {pixel_values.shape}")
         
         latents = pipeline.vae.apply(