Hi! I tried using nn.Linear module and tried to flatten the last layer of my network too. My code finally has no error but still I am not 100% sure if I did it correctly, any comments on this sir? :
class ESNet(nn.Module):
    def __init__(self, classes):
        super().__init__()
        #-----ESNET---------#
        self.initial_block = DownsamplerBlock(3, 16)

        self.layers = nn.ModuleList()
        
        for x in range(0, 3):
           self.layers.append(FCU(16, 3, 0.03, 1))  
        
        self.layers.append(DownsamplerBlock(16,64))

        for x in range(0, 2):
           self.layers.append(FCU(64, 5, 0.03, 1))  

        self.layers.append(DownsamplerBlock(64,128))

        for x in range(0, 3):   
            self.layers.append(PFCU(chann=128)) 

        self.layers.append(UpsamplerBlock(128,64))
        self.layers.append(FCU(64, 5, 0, 1))
        self.layers.append(FCU(64, 5, 0, 1))

        self.layers.append(UpsamplerBlock(64,16))
        self.layers.append(FCU(16, 3, 0, 1))
        self.layers.append(FCU(16, 3, 0, 1))

        self.output_conv = nn.ConvTranspose2d( 16, classes, 2, stride=2, padding=0, output_padding=0, bias=True)
        self.hidden = nn.Linear(2*104*104, 208)
        self.out = nn.Linear(208, 2)
        self.act = nn.ReLU()

    def forward(self, input):
        output = self.initial_block(input)
        print(input.shape)
        

        for layer in self.layers:
            output = layer(output)
           

        output = self.output_conv(output)
        output = output.view(output.size(0), -1)
        output = self.act(self.hidden(output))
        output = self.out(output)
        print(output.shape)
       
        return output