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

Enable find_unused_parameters for DistributedDataParallel

Browse Source 
attention_norm has some parameters which are not used to compute grad,
which is causing failures in the distributed case.
This commit is contained in:
James Betker 2020-07-23 09:08:13 -06:00
parent dbf6147504
commit bba283776c
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
codes/models/SRGAN_model.py
View File

@ -136,13 +136,15 @@ class SRGANModel(BaseModel):
# DataParallel
if opt['dist']:
self.netG = DistributedDataParallel(self.netG, device_ids=[torch.cuda.current_device()])
self.netG = DistributedDataParallel(self.netG, device_ids=[torch.cuda.current_device()],
find_unused_parameters=True)
else:
self.netG = DataParallel(self.netG)
if self.is_train:
if opt['dist']:
self.netD = DistributedDataParallel(self.netD,
device_ids=[torch.cuda.current_device()])
device_ids=[torch.cuda.current_device()],
find_unused_parameters=True)
else:
self.netD = DataParallel(self.netD)
self.netG.train()