Update app.py

app.py

@@ -161,4 +161,28 @@ state = train_state.TrainState.create(
 # Training loop
 num_epochs = 3
 batch_size = 1
-rng = jax.random.PRNGKey(
+rng = jax.random.PRNGKey(0)
+
+for epoch in range(num_epochs):
+    epoch_loss = 0
+    num_batches = 0
+    for batch in tqdm(processed_dataset.batch(batch_size)):
+        batch['pixel_values'] = jnp.array(batch['pixel_values'], 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()
+    
+    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)
+unet.save_pretrained(output_dir, params=state.params["unet"])
+
+print(f"Model saved to {output_dir}")