Explorar o código

Update train.py

Peter Lin %!s(int64=3) %!d(string=hai) anos
pai
achega
48effc9157
Modificáronse 1 ficheiros con 2 adicións e 0 borrados
  1. 2 0
      train.py

+ 2 - 0
train.py

@@ -329,6 +329,8 @@ class Trainer:
             {'params': self.model.backbone.parameters(), 'lr': self.args.learning_rate_backbone},
             {'params': self.model.aspp.parameters(), 'lr': self.args.learning_rate_aspp},
             {'params': self.model.decoder.parameters(), 'lr': self.args.learning_rate_decoder},
+            {'params': self.model.project_mat.parameters(), 'lr': self.args.learning_rate_decoder},
+            {'params': self.model.project_seg.parameters(), 'lr': self.args.learning_rate_decoder},
             {'params': self.model.refiner.parameters(), 'lr': self.args.learning_rate_refiner},
         ])
         self.scaler = GradScaler()