Update app.py

app.py

@@ -32,11 +32,9 @@ def start_train(status_box):
     # Handles CUDA OOM errors.
     os.system(f"export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True")
     os.system("echo 'Okay, trying training.'")
-    status_box.value = "Status: Training"
     API.add_space_variable(repo_id=experiment_repo, key="STATUS", value="TRAINING")
     os.system(f"cd pytorch-image-models; ./train.sh 4 --dataset hfds/datacomp/imagenet-1k-random-10.0-frac-1over4 --log-wandb --wandb-project ImageNetTraining10.0-frac-1over4 --experiment ImageNetTraining10.0-frac-1over4 --model seresnet34 --sched cosine --epochs 150 --warmup-epochs 5 --lr 0.4 --reprob 0.5 --remode pixel --batch-size 256 --amp -j 4")
     API.add_space_variable(repo_id=experiment_repo, key="STATUS", value="IDLE")
-    status_box.value = "Status: Idle"
     os.system("echo 'Done'.")
     os.system("ls")
     # Upload output to repository
@@ -54,7 +52,7 @@ def run():
             API.add_space_variable(repo_id=experiment_repo, key="STATUS", value="IDLE")
             status_box.value = "Status: Idle"
             start = gr.Button("Start")
-            start.click(start_train, inputs=status_box)
+            start.click(start_train, inputs=status_box, outputs=status_box)
         else:
             status_box.value = "Status: Training"
     app.launch(server_name="0.0.0.0", server_port=7860)