try again..

This commit is contained in:
James Betker 2020-09-22 18:27:52 -06:00
parent 3b6c957194
commit 4ab989e015

View File

@ -82,7 +82,7 @@ class FeatureLoss(ConfigurableLoss):
logits_real = self.netF(state[self.opt['real']])
logits_fake = self.netF(state[self.opt['fake']])
if self.opt['criterion'] == 'cosine':
return self.criterion(logits_fake, logits_real, torch.ones_like(1, device=logits_fake.device))
return self.criterion(logits_fake, logits_real, torch.ones(1, device=logits_fake.device))
else:
return self.criterion(logits_fake, logits_real)