Update app.py

app.py

@@ -238,26 +238,42 @@ num_epochs = 3
 batch_size = 1
 rng = jax.random.PRNGKey(0)
 
+# Training loop
+num_epochs = 3
+batch_size = 1
+rng = jax.random.PRNGKey(0)
+
 for epoch in range(num_epochs):
     epoch_loss = 0
     num_batches = 0
+    num_errors = 0
     for batch in tqdm(processed_dataset.batch(batch_size)):
-        batch['pixel_values'] = jnp.array(batch['pixel_values'][0], dtype=jnp.float32)
-        rng, step_rng = jax.random.split(rng)
-        state, loss = train_step(state, batch, step_rng)
-        epoch_loss += loss
-        num_batches += 1
-        
-        if num_batches % 10 == 0:
-            jax.clear_caches()
+        try:
+            batch['pixel_values'] = jnp.array(batch['pixel_values'][0], dtype=jnp.float32)
+            rng, step_rng = jax.random.split(rng)
+            state, loss = train_step(state, batch, step_rng)
+            epoch_loss += loss
+            num_batches += 1
+            
+            if num_batches % 10 == 0:
+                jax.clear_caches()
+                print(f"Processed {num_batches} batches. Current loss: {loss}")
+        except Exception as e:
+            num_errors += 1
+            print(f"Error processing batch: {e}")
+            continue
+    
+    if num_batches > 0:
+        avg_loss = epoch_loss / num_batches
+        print(f"Epoch {epoch+1}/{num_epochs}, Average Loss: {avg_loss}, Errors: {num_errors}")
+    else:
+        print(f"Epoch {epoch+1}/{num_epochs}, No valid batches processed, Errors: {num_errors}")
     
-    avg_loss = epoch_loss / num_batches
-    print(f"Epoch {epoch+1}/{num_epochs}, Average Loss: {avg_loss}")
     jax.clear_caches()
 
 # Save the fine-tuned model
 output_dir = "/tmp/montevideo_fine_tuned_model"
 os.makedirs(output_dir, exist_ok=True)
-adjusted_unet.save_pretrained(output_dir, params=state.params)
+adjusted_unet.save_pretrained(output_dir, params=state.params["params"])
 
 print(f"Model saved to {output_dir}")