bug re-fix

thuml · Apr 28, 2020 · 5d7caa9 · 5d7caa9
1 parent 4422564
commit 5d7caa9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/main.py b/main.py
@@ -181,8 +181,8 @@ def outlier(each_target_share_weight):
         adv_loss_separate += nn.BCELoss()(domain_prob_discriminator_source_separate, torch.ones_like(domain_prob_discriminator_source_separate))
         adv_loss_separate += nn.BCELoss()(domain_prob_discriminator_target_separate, torch.zeros_like(domain_prob_discriminator_target_separate))
 
-        # ============================== cross entropy loss, it receives logits as its inputs
-        ce = nn.CrossEntropyLoss(reduction='none')(fc2_s, label_source)
+        # ============================== cross entropy loss
+        ce = nn.CrossEntropyLoss(reduction='none')(predict_prob_source, label_source)
         ce = torch.mean(ce, dim=0, keepdim=True)
 
         with OptimizerManager(