kxngluca/DL-Art-School
1
0
Fork 0
forked from mrq/DL-Art-School
Code Pull Requests Activity

Allow swapping to torch DDP as needed in code

Browse Source
This commit is contained in:
James Betker 2020-12-09 15:03:59 -07:00
parent 66cbae8731
commit c203cee31e
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
codes/models/ExtensibleTrainer.py
View File

@ -108,6 +108,8 @@ class ExtensibleTrainer(BaseModel):
# Use Apex to enable delay_allreduce, which is compatible with gradient checkpointing.
from apex.parallel import DistributedDataParallel
dnet = DistributedDataParallel(anet, delay_allreduce=True)
#from torch.nn.parallel.distributed import DistributedDataParallel
#dnet = DistributedDataParallel(anet, device_ids=[torch.cuda.current_device()], find_unused_parameters=True)
else:
dnet = DataParallel(anet, device_ids=opt['gpu_ids'])
if self.is_train: