1

Worromots · May 23, 2022 · 3d42692 · 3d42692
1 parent d965bd8
commit 3d42692
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/.idea/vcs.xml b/.idea/vcs.xml
diff --git a/train.py b/train.py
@@ -64,10 +64,12 @@ def _get_model(self):
         return tfnet
 
     def _get_optimizer(self, model):
-        return torch.optim.SGD([
-            {"params": model.patten.parameters(), "lr": self.cfg.learn_rate_seg}
-        ],
-            lr=self.cfg.learn_rate)
+        # return torch.optim.SGD([
+        #     {"params": model.patten.parameters(), "lr": self.cfg.learn_rate_seg}
+        # ],
+        #     lr=self.cfg.learn_rate)
+        return torch.optim.SGD(model.parameters(), self.cfg.LEARNING_RATE)
+
 
     def _train_model(self, device, model, train_loader, criterion_seg, criterion_dec, optimizer, test_loader,
                      tensorboard_writer):