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

Fix another issue with gpu ids getting thrown all over hte place

Browse Source
This commit is contained in:
James Betker 2020-11-13 20:05:52 -07:00
parent c47925ae34
commit 0b96811611
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
codes/models/steps/losses.py
View File

@ -134,7 +134,7 @@ class FeatureLoss(ConfigurableLoss):
self.netF = define_F(which_model=opt['which_model_F'],
load_path=opt['load_path'] if 'load_path' in opt.keys() else None).to(self.env['device'])
if not env['opt']['dist']:
self.netF = torch.nn.parallel.DataParallel(self.netF)
self.netF = torch.nn.parallel.DataParallel(self.netF, device_ids=env['opt']['gpu_ids'])
def forward(self, _, state):
with autocast(enabled=self.env['opt']['fp16']):