Is flattening_layer correct ?

[chrisM38]

I have a question. I was trying to train the colorization_cINN, but I have some struggles with the fully connected section. I keep getting an error in the glow_coupling_layer:

File "...\FrEIA\modules\coupling_layers.py", line 179, in forward
    r2 = self.s2(torch.cat([x2, *c], 1) if self.conditional else x2)
RuntimeError: Tensors must have same number of dimensions: got 1 and 2

X2 keeps having 2 Dimensions, even after applying the flattening_layer. The output_shape is giving you the correct format.

x = torch.rand((1, 32, 4, 4))

module = flattening_layer ([x.size()])
print(module.output_shape)   ## (512,)
print(module.forward([x])[0].size()) ## torch.Size([1, 512])

Am I doing something wrong there ?
Thanks in advance!