fix interface bug

Author: SanghyukChun

criterions/comparison_methods.py

@@ -80,7 +80,7 @@ def __init__(self, w=0.36, **kwargs):
         self.w = w
         self.fc = nn.Linear(kwargs.get('feat_dim', 128), 1).cuda()
 
-    def forward(self, f_feat, g_pred, labels, f_pred):
+    def forward(self, f_feat, g_feat, labels, f_pred, g_pred):
         f_feat = f_feat.view(f_feat.shape[0], -1)
         f_pred = f_pred.view(f_pred.shape[0], -1)
         g_pred = g_pred.view(g_pred.shape[0], -1)
@@ -89,7 +89,7 @@ def forward(self, f_feat, g_pred, labels, f_pred):
         factor = F.softplus(factor)
         g_pred *= factor
 
-        loss = F.cross_entropy(f_pred+g_pred, labels)
+        loss = F.cross_entropy(f_pred + g_pred, labels)
 
         bias_lp = F.log_softmax(g_pred, 1)
         entropy = -(torch.exp(bias_lp) * bias_lp).sum(1).mean()

trainer.py

@@ -337,7 +337,7 @@ def _update_f(self, x, labels, update_outer_loop=True, loss_dict=None, prefix=''
             for g_idx, g_net in enumerate(self.model.g_nets):
                 _g_preds, _g_feats = g_net(x)
 
-                _f_loss_indep = self.outer_criterion(f_feats, _g_feats, labels=labels, f_pred=preds)
+                _f_loss_indep = self.outer_criterion(f_feats, _g_feats, labels=labels, f_pred=preds, g_pred=_g_preds)
                 f_loss_indep += _f_loss_indep
 
                 loss_dict['{}f_loss_indep_g_{}'.format(prefix, g_idx)] = _f_loss_indep.item()