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

Fix distributed d_grad

Browse Source
This commit is contained in:
James Betker 2020-08-25 16:06:27 -06:00
parent 09a9079e17
commit 96586d6592
1 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
codes/models/SRGAN_model.py
View File

@ -288,8 +288,14 @@ class SRGANModel(BaseModel):
self.netD = DistributedDataParallel(self.netD,
device_ids=[torch.cuda.current_device()],
find_unused_parameters=True)
if self.spsr_enabled:
self.netD_grad = DistributedDataParallel(self.netD_grad,
device_ids=[torch.cuda.current_device()],
find_unused_parameters=True)
else:
self.netD = DataParallel(self.netD)
if self.spsr_enabled:
self.netD_grad = DataParallel(self.netD_grad)
self.netG.train()
self.netD.train()
if self.spsr_enabled: