deviandice/DL-Art-School
1
0
Fork 0
forked from mrq/DL-Art-School
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix overuse of checkpointing

Browse Source
This commit is contained in:
James Betker 2021-09-16 23:00:28 -06:00
parent f78ce9d924
commit 94899d88f3
2 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
codes/models/diffusion/diffusion_dvae.py
View File

@ -32,7 +32,7 @@ class DiscreteEncoder(nn.Module):
) )
def forward(self, spectrogram): def forward(self, spectrogram):
return checkpoint(self.blocks, spectrogram) return self.blocks(spectrogram)
class DiscreteDecoder(nn.Module): class DiscreteDecoder(nn.Module):

2
codes/trainer/ExtensibleTrainer.py
View File

@ -117,7 +117,7 @@ class ExtensibleTrainer(BaseModel):
dnet = DistributedDataParallel(anet, delay_allreduce=True) dnet = DistributedDataParallel(anet, delay_allreduce=True)
else: else:
from torch.nn.parallel.distributed import DistributedDataParallel from torch.nn.parallel.distributed import DistributedDataParallel
dnet = DistributedDataParallel(anet, device_ids=[torch.cuda.current_device()]) dnet = DistributedDataParallel(anet, device_ids=[torch.cuda.current_device()], find_unused_parameters=True)
else: else:
dnet = DataParallel(anet, device_ids=opt['gpu_ids']) dnet = DataParallel(anet, device_ids=opt['gpu_ids'])
if self.is_train: if self.is_train: