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

add checkpointing

Browse Source
This commit is contained in:
James Betker 2022-05-31 21:09:05 -06:00
parent c0db85bf4f
commit 8a1b8e3e62
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
codes/models/audio/music/music_quantizer.py
View File

@ -4,6 +4,7 @@ import torch.nn.functional as F
from models.arch_util import zero_module
from trainer.networks import register_model
from utils.util import checkpoint
class Downsample(nn.Module):
@ -42,7 +43,10 @@ class ResBlock(nn.Module):
)
def forward(self, x):
return self.net(x) + x
return checkpoint(self._forward, x) + x
def _forward(self, x):
return self.net(x)
class Wav2Vec2GumbelVectorQuantizer(nn.Module):