How can I add reshape layer in nn.Sequential?
So I'm implementing Generator of a GAN and I need the architecture as shown as below:

The problem is when I try to reshape the output of Linear layer after BatchNorm and ReLU (in fig. Dense as they have used Tensorflow) it is throwing error as :TypeError: reshape(): argument 'input' (position 1) must be Tensor, not int
I understand the error but I can't find its solution.
Is there any other way to reshape within nn.Sequential instead of calling torch explicitly?
class Generator(nn.Module):

    def __init__(self, z_dim=100, im_chan=1, hidden_dim=64, rdim=9216):
       
        super(Generator, self).__init__()
        self.z_dim = z_dim
        self.gen = nn.Sequential(
             nn.Linear(z_dim, rdim),
             nn.BatchNorm2d(rdim,momentum=0.9),
             nn.ReLU(inplace=True),
 ----&gt;       torch.reshape(rdim, (6,6,256)), 
             self.make_gen_block(rdim, hidden_dim*2),
             self.make_gen_block(hidden_dim*2,hidden_dim),
             self.make_gen_block(hidden_dim,im_chan,final_layer=True),
        )
   def make_gen_block(self, input_channels, output_channels, kernel_size=1, stride=2, final_layer=False):

        if not final_layer:
           return nn.Sequential(
              nn.ConvTranspose2d(input_channels, output_channels, kernel_size, stride),
              nn.BatchNorm2d(output_channels),
              nn.ReLU(inplace=True)
        )
        else:
              return nn.Sequential(
              nn.ConvTranspose2d(input_channels, output_channels, kernel_size, stride),
              nn.Tanh()
        )

   def unsqueeze_noise(self, noise):
       return noise.view(len(noise), self.zdim, 1, 1)

   def forward(self, noise):
       x = self.unsqueeze_noise(noise)
       return self.gen(x)

def get_noise(n_samples, z_dim, device='cpu'):
    return torch.randn(n_samples, z_dim, device=device)

#Testing the Gen arch
gen = Generator()
num_test = 100

#test the hidden block
test_hidden_noise = get_noise(num_test, gen.z_dim)
test_hidden_block = gen.make_gen_block(6, 6, kernel_size=1,stride=2)
test_uns_noise = gen.unsqueeze_noise(test_hidden_noise)
hidden_output = test_hidden_block(test_uns_noise)



In nn.Sequential, torch.nn.Unflatten() can help you achieve reshape operation.

For nn.Linear, its input shape is (N, *, H_{in}) and output shape is (H, *, H_{out}). Note that the feature dimension is last. So unsqueeze_noise() is not useful here.

Based on the network structure, the arguments passed to make_gen_block are wrong.


I have checked the following code:
import torch
from torch import nn
class Generator(nn.Module):

   def __init__(self, z_dim=100, im_chan=1, hidden_dim=64, rdim=9216):
       
        super(Generator, self).__init__()
        self.z_dim = z_dim
        self.gen = nn.Sequential(
             nn.Linear(z_dim, rdim),
             nn.BatchNorm1d(rdim,momentum=0.9), # use BN1d
             nn.ReLU(inplace=True),
             nn.Unflatten(1, (256,6,6)), 
             self.make_gen_block(256, hidden_dim*2,kernel_size=2), # note arguments
             self.make_gen_block(hidden_dim*2,hidden_dim,kernel_size=2), # note kernel_size
             self.make_gen_block(hidden_dim,im_chan,kernel_size=2,final_layer=True), # note kernel_size
        )
   def make_gen_block(self, input_channels, output_channels, kernel_size=1, stride=2, final_layer=False):

        if not final_layer:
           return nn.Sequential(
              nn.ConvTranspose2d(input_channels, output_channels, kernel_size, stride),
              nn.BatchNorm2d(output_channels),
              nn.ReLU(inplace=True)
        )
        else:
              return nn.Sequential(
              nn.ConvTranspose2d(input_channels, output_channels, kernel_size, stride),
              nn.Tanh()
        )

   def forward(self, x):
       return self.gen(x)

def get_noise(n_samples, z_dim, device='cpu'):
    return torch.randn(n_samples, z_dim, device=device)

gen = Generator()
num_test = 100
input_noise = get_noise(num_test, gen.z_dim)
output = gen(input_noise)
assert output.shape == (num_test, 1, 48, 48)

