Solved by ZimoNitrome in post #2 
              
              
                I am not exactly sure what you mean. But if you want to debug the shape (and hidden dim) of the tensor between each block, I like to use a simple “PrintBlock”: 
class PrintBlock(nn.Module):
    def forward(self, x):
        print(x.shape)
        return x

Putting it into your code: 
class Generato…