Move the model to GPU before constructing optimizer

1 files changed, 7 insertions, 6 deletions

diff --git a/models/model.py b/models/model.py
index e8d79f9..3cae788 100644
--- a/models/model.py
+++ b/models/model.py
@@ -115,11 +115,11 @@ class Model:
         optim_hp = self.hp.get('optimizer', {})
         sched_hp = self.hp.get('scheduler', {})
         self.rgb_pn = RGBPartNet(self.train_size, self.in_channels, **model_hp)
+        # Try to accelerate computation using CUDA or others
+        self.rgb_pn = self._accelerate(self.rgb_pn)
         self.optimizer = optim.Adam(self.rgb_pn.parameters(), **optim_hp)
         self.scheduler = optim.lr_scheduler.StepLR(self.optimizer, **sched_hp)
         self.writer = SummaryWriter(self._log_name)
-        # Try to accelerate computation using CUDA or others
-        self._accelerate()
 
         self.rgb_pn.train()
         # Init weights at first iter
@@ -176,11 +176,12 @@ class Model:
                 self.writer.close()
                 break
 
-    def _accelerate(self):
+    def _accelerate(self, model: nn.Module) -> nn.Module:
         if not self.disable_acc:
             if torch.cuda.device_count() > 1:
-                self.rgb_pn = nn.DataParallel(self.rgb_pn)
-            self.rgb_pn = self.rgb_pn.to(self.device)
+                model = nn.DataParallel(model)
+            model = model.to(self.device)
+        return model
 
     def predict_all(
             self,
@@ -204,7 +205,7 @@ class Model:
         model_hp = self.hp.get('model', {})
         self.rgb_pn = RGBPartNet(ae_in_channels=self.in_channels, **model_hp)
         # Try to accelerate computation using CUDA or others
-        self._accelerate()
+        self.rgb_pn = self._accelerate(self.rgb_pn)
 
         self.rgb_pn.eval()
         gallery_samples, probe_samples = [], {}