Update app.py

app.py

@@ -223,12 +223,7 @@ def log_message(output_log, msg):
 # =====================================================
 @spaces.GPU(duration=300)
 def train_model(
-    base_model: str,
-    dataset_name: str,
-    hf_repo: str,
-    num_epochs: int = 1,
-    batch_size: int = 2,
-    learning_rate: float = 5e-4,
+   base_model, dataset_name, num_epochs, batch_size, learning_rate, hf_repo
 ):
     output_log = []