extension

[jecampagne]

Very nice job! One little extension concerns the number of channels, ie. RGB=3 but it can nicely be genralized

def init(self, num_layers=5, num_features=64, num_channels=1):

then the 1st conv_layer

nn.Conv2d(num_channels, ... and the last deconv_layer deconv_layers.append(nn.ConvTranspose2d(num_features, num_channels,... Thanks