mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix feature validation, wrong device

Browse Source 
Only shows up in distributed training for some reason.
This commit is contained in:
James Betker 2020-07-23 10:16:34 -06:00
parent a7541b6d8d
commit 9ccf771629
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
codes/models/SRGAN_model.py
View File

@ -554,8 +554,8 @@ class SRGANModel(BaseModel):
def compute_fea_loss(self, real, fake):
with torch.no_grad():
real = real.unsqueeze(dim=0)
fake = fake.unsqueeze(dim=0)
real = real.unsqueeze(dim=0).to(self.device)
fake = fake.unsqueeze(dim=0).to(self.device)
real_fea = self.netF(real).detach()
fake_fea = self.netF(fake)
return self.cri_fea(fake_fea, real_fea).item()