Reshaping real/imaginary parts into channel dimension problem in NCSNpp class forward method

[mattpitkin]

In the forward method of the NCSNpp class, the reshaping of the input to put real and imaginary parts into the channels dimension works if spatial_dimensions is 1, but if it were greater than 1 it does not work properly. I think to following:

x_chans = []
for chan in range(self.spatial_channels):
    x_chans.append(torch.cat([ 
        torch.cat([x[:,[chan+in_chan],:,:].real, x[:,[chan+in_chan],:,:].imag ], dim=1) for in_chan in range(self.input_channels // 2)],
            dim=1)
    )
x = torch.cat(x_chans, dim=1) #4*D

needs to be replaced with:

x_chans = []
for chan in range(self.spatial_channels):
    x_chans.append(torch.cat([ 
        torch.cat([x[:,[in_chan + (chan * input_channels // 2)],:,:].real, x[:,[in_chan + (chan * input_channels // 2)],:,:].imag ], dim=1) for in_chan in range(self.input_channels // 2)],
            dim=1)
    )
x = torch.cat(x_chans, dim=1) #4*D

or more compactly, just with:

x = torch.view_as_real(x).permute(0, 1, 4, 2, 3).flatten(1, 2)