@omarfoq
Below is a part of the code for my model, disc.
def forward(self, inputs):
................
..... More code
...............
        x3 = x2.max(dim=-1, keepdim=False)[0]
        em = x3
        
        x = F.leaky_relu(self.bn6(self.linear1(x3)), negative_slope=0.2) # (batch_size, emb_dims*2) -> (batch_size, 512)
        #x = self.dp1(x)
        x = F.leaky_relu(self.bn7(self.linear2(x)), negative_slope=0.2) # (batch_size, 512) -> (batch_size, 256)
        x = self.dp2(x)
        x = self.linear3(x)                                             # (batch_size, 256) -> (batch_size, output_channels)
        
        return x,em

em is just the output striaght from the convolution and maxpooling.