Merge pull request #69 from ikvision/patch-3

Update trainer.py
nkolot · Apr 1, 2020 · 5c79685 · 5c79685
2 parents 21a148d + d7dbd81
commit 5c79685
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train/trainer.py b/train/trainer.py
@@ -255,7 +255,7 @@ def train_step(self, input_batch):
         loss = self.options.shape_loss_weight * loss_shape +\
                self.options.keypoint_loss_weight * loss_keypoints +\
                self.options.keypoint_loss_weight * loss_keypoints_3d +\
-               loss_regr_pose + self.options.beta_loss_weight * loss_regr_betas +\
+               self.options.pose_loss_weight * loss_regr_pose + self.options.beta_loss_weight * loss_regr_betas +\
                ((torch.exp(-pred_camera[:,0]*10)) ** 2 ).mean()
         loss *= 60