Update app.py

app.py

@@ -95,6 +95,10 @@ def train_lora(image_folder, metadata):
             if batch_idx % 10 == 0:
                 print(f"Batch {batch_idx}, Loss: {loss.item()}")
 
+# Save the trained model
+torch.save(model.state_dict(), "lora_model.pth")
+print("Model saved as lora_model.pth")
+    
     print("LoRA training completed.")
 
 # Gradio App