diff --git a/workspace/odelia-breast-mri/model/main.py b/workspace/odelia-breast-mri/model/main.py
index 50534e3e..969d7329 100755
--- a/workspace/odelia-breast-mri/model/main.py
+++ b/workspace/odelia-breast-mri/model/main.py
@@ -223,25 +223,7 @@ def on_validation_epoch_end(self, trainer, pl_module):
         swarmCallback.on_train_end()
     model.save_best_checkpoint(trainer.logger.log_dir, checkpointing.best_model_path)
     model.save_last_checkpoint(trainer.logger.log_dir, checkpointing.last_model_path)
-    '''
-    import subprocess
 
-    # Get the container ID for the latest user-env container
-    get_container_id_command = 'docker ps -a --filter "name=us*" --format "{{.ID}}" | head -n 1'
-    container_id = subprocess.check_output(get_container_id_command, shell=True, text=True).strip()
-
-    # Get the latest log for the user-env container
-    get_logs_command = f"docker logs {container_id}"
-    logs_process = subprocess.Popen(get_logs_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-
-    # Print and log the output
-    with open(os.path.join(path_run_dir,"container_logs.txt"), "w") as log_file:
-        for line in logs_process.stdout:
-            line = line.decode("utf-8").rstrip()
-            print(line)
-            log_file.write(line + "\n")
-    '''
 
     predict(path_run_dir, os.path.join(dataDir, task_data_name,'test'), model_name, last_flag=False, prediction_flag = prediction_flag)
-    predict(path_run_dir, os.path.join(dataDir, task_data_name,'test'), model_name, last_flag=True, prediction_flag = prediction_flag)
-
+    print('Training completed successfully')