fix inconsistent in loss function and training

Andrew05200 · Aug 14, 2019 · e2786c5 · e2786c5
1 parent 7a1b216
commit e2786c5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/model/loss.py b/model/loss.py
@@ -134,10 +134,10 @@ def forward(self, logits, target,
         else:
             if is_multi:
                 assert self.loss_type in [LossType.BCE_WITH_LOGITS,
-                                        LossType.SIGMOID_FOCAL_CROSS_ENTROPY]
+                                          LossType.SIGMOID_FOCAL_CROSS_ENTROPY]
             else:
                 if self.loss_type not in [LossType.SOFTMAX_CROSS_ENTROPY,
-                                          LossType.SOFTMAX_FOCAL_CROSS_ENTROPY]
+                                          LossType.SOFTMAX_FOCAL_CROSS_ENTROPY]:
                     target = torch.eye(self.label_size)[target].to(device)
             return self.criterion(logits, target)