Unet1dmodel fails

Specifying the architecture of a Unet1dmodel leads to weird behaviour.
CODE :

default_model = UNet1DModel(in_channels = 17, out_channels=1,sample_size = 256)
custom_model = UNet1DModel(in_channels = 17, out_channels=1,sample_size = 256,down_block_types =
(โ€œDownBlock1Dโ€, โ€œDownBlock1Dโ€, โ€œDownBlock1Dโ€,โ€œDownBlock1Dโ€,โ€œDownBlock1Dโ€),
up_block_types = (โ€œUpBlock1Dโ€, โ€œUpBlock1Dโ€,โ€œUpBlock1Dโ€,โ€œUpBlock1Dโ€, โ€œUpBlock1Dโ€),
block_out_channels = (32, 32, 32, 32, 32))
x = torch.randn((16,1,256))
timesteps = torch.randint(0,10,(16,))
y_default = default_model(x,timesteps)
y_custom = custom_model(x,timesteps)

ERROR :
Regarding y_custom
RuntimeError: Given groups=1, weight of size [32, 17, 1], expected input[16, 1, 128] to have 17 channels, but got 1 channels instead

The in_channel argument is simply due to the following thread

hey @StableDiffuser317 Iโ€™m not sure how much time we have to prioritize debugging the 1 dimensional unet beyond the default configuration but please feel free to do some digging and open a PR updating the misconfiguration :slight_smile: