Update app.py

app.py

@@ -188,13 +188,20 @@ def train_model(epochs):
             loss.backward()
             optimizer.step()
 
-            if i % 100 == 0:
+            if i % 10 == 0:
                 status = f"Epoch [{epoch}/{epochs}], Step [{i}/{len(dataloader)}], Loss: {loss.item():.8f}"
                 print(status)
                 output_text.append(status)
 
     return model, "\n".join(output_text)
 
+# Push model to Hugging Face Hub
+def push_model_to_hub(model, repo_id):
+    wrapper = UNetWrapper(model, repo_id)
+    wrapper.push_to_hub()
+    # Push the model to the Hugging Face hub
+    #model.push_to_hub(repo_name)
+    
 # Gradio interface function
 def gradio_train(epochs):
     model, training_log = train_model(int(epochs))