From 6cf203e173c13a031a02e1a49640ae66a881a066 Mon Sep 17 00:00:00 2001
From: um3 <zhedongzheng@um.edu.mo>
Date: Wed, 24 Jan 2024 18:05:26 +0800
Subject: [PATCH] update DDP last

---
 train_DDP.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/train_DDP.py b/train_DDP.py
index 105e410..c3d20d3 100755
--- a/train_DDP.py
+++ b/train_DDP.py
@@ -451,10 +451,7 @@ def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
 
     # load best model weights
     model.load_state_dict(last_model_wts)
-    if len(opt.gpu_ids)>1:
-        save_network(model.module, 'last')
-    else:
-        save_network(model, 'last')
+    save_network(model.module, opt.name, 'last')
 
     return model