Solved by ptrblck in post #4 
              
              
                Thanks for the update! 
There seems to be a small misunderstanding. 
You should still create the tensors in the default format [N, C, H, W] and just call to(memory_format=torch.channels_last) on it, so that your code changes would be minimal. 
This code works for me: 
model = Generator((2048, 1, 2))…