def __init__(self):
super(WaveVAE, self).__init__()
self.encoder = Wavenet(out_channels=2,
num_blocks=2,
num_layers=10,
residual_channels=128,
gate_channels=256,
skip_channels=128,
kernel_size=2,
cin_channels=80,
upsample_scales=[16, 16])
self.decoder = Wavenet_Student(num_blocks_student=[1, 1, 1, 1, 1, 1],
num_layers=10)
self.log_eps = nn.Parameter(torch.zeros(1))
def build_student():
model_s = Wavenet_Student(num_blocks_student=[1, 1, 1, 1, 1, 1],
num_layers=args.num_layers_s)
return model_s