deviandice/DL-Art-School
1
0
Fork 0
forked from mrq/DL-Art-School
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()