Wrong dilation order

[yukara-ikemiya]

Hi, thank you for the implementation of WaveGrad. It helps me a lot for implementing WaveFit from Google in the following repository.
https://github.com/yukara-ikemiya/wavefit-pytorch

I happened to find the dilation order of upsampling blocks might be wrong.
The original paper says the dilation factor 1, 2, 4, 8 should be used for the first 3 blocks, but in your implementation, it is used for the last 3 blocks.

I believe this should be modified as below.

wavegrad/src/wavegrad/model.py

Lines 169 to 175 in c3d54c8

	self.upsample = nn.ModuleList([
	UBlock(768, 512, 5, [1, 2, 1, 2]),
	UBlock(512, 512, 5, [1, 2, 1, 2]),
	UBlock(512, 256, 3, [1, 2, 4, 8]),
	UBlock(256, 128, 2, [1, 2, 4, 8]),
	UBlock(128, 128, 2, [1, 2, 4, 8]),
	])

self.upsample = nn.ModuleList([
        UBlock(768, 512, 5, [1, 2, 4, 8]),
        UBlock(512, 512, 5, [1, 2, 4, 8]),
        UBlock(512, 256, 3, [1, 2, 4, 8]),
        UBlock(256, 128, 2, [1, 2, 1, 2]),
        UBlock(128, 128, 2, [1, 2, 1, 2]),
    ])

I conducted small scale experiments and confirmed that this modification actually improved the performance.