Final ResNetBlock block for UNet in denoising_diffusion_pytorch.py

[MekkCyber]

Hello @lucidrains

Sorry if i'm mistaken, but in the line 380 of the file denoising_diffusion_pytorch.py, where the final ResNetBlock is defined for the UNet :
self.final_res_block = block_klass(dim * 2, dim, time_emb_dim = time_dim)
Shouldn't the input channels dimension be init_dim * 2 instead of dim * 2 ?