added timeout handling for docker container removal

run.sh

@@ -14,7 +14,12 @@ TRAIN_ARGS="--save_dir ./ReAct/outputs/ --dataset 'minipile' --group 'minipile'
 
 # Stop all running Docker containers
 echo "Stopping all running Docker containers..."
-sudo docker rm -f $CONTAINER_NAME
+
+if ! timeout 300 sudo docker rm -f $CONTAINER_NAME; then
+    echo "Command timed out. Restarting Docker daemon & retrying..."
+    sudo systemctl restart docker
+    sleep 10s; sudo docker rm -f $CONTAINER_NAME
+fi
 
 # Git stuff
 git clone -b $BRANCH https://github.com/neel04/ReAct_Jax.git

train_model.py

@@ -91,7 +91,7 @@ def main(key: PRNGKeyArray):
         )
 
         # enqueue a few handpicked hyperparams for trials
-        _: None = [study.enqueue_trial(hyperparams) for hyperparams in init_hyperparams]
+        [study.enqueue_trial(hyperparams) for hyperparams in init_hyperparams]
 
         study.optimize(
             lambda trial: kickoff_optuna(trial=trial, **trainer_kwargs),