Thank you very much for answering!
I wanted to implement your suggestion of reshaping the activation back but I couldn’t figure out how to go from 2-dimensional weight [2,1280] to 4-dimensional [1280, 1280, 2, 2] (I guess the 2,2 are due to the kernel and stride).
So, I tried to define the forward only on the mobilenetv2.features as written below (thought this way I will just avoid the activation function) but that resulted in the next error “Expected 4-dimensional input for 4-dimensional weight [1280, 1280, 1, 1], but got 2-dimensional input of size [2, 62720] instead”.
Here’s the current code:
class Flatten(nn.Module):
def forward(self, input):
return input.view(input.size(0), -1)
class MobileNet_ConvAdd(nn.Module):
def init(self):
super(MobileNet_ConvAdd, self).init()
self.model_M = models.mobilenet_v2(pretrained=True)
self.model_M.classifier = nn.Identity()
for param in self.model_M.parameters():
param.requires_grad = False
    self.MobileNet_ConvAdd_conv1 = nn.Sequential(
                                                nn.Flatten(),
                                                nn.Conv2d(1280, 1280, kernel_size=(2, 2), stride=(2, 2), bias=False),
                                                nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
                                                nn.ReLU(),
                                                nn.Conv1d(1280, 320, kernel_size=(1, 1), stride=(1, 1), bias=False),
                                                nn.AvgPool2d(kernel_size=2, stride=2, padding=0),
                                                nn.Flatten()
                                                )

    self.MobileNet_ConvAdd_fc2 = nn.Sequential(
                                    nn.Dropout(0.2),
                                    nn.BatchNorm1d(1280), #320
                                    nn.ReLU(),
                                    nn.Dropout(0.3),      
                                    nn.Linear(320,NUM_CLASSES),
                                    nn.LogSoftmax(dim=1)
                                    )

def forward(self, x):
    x = self.model_M.features(x)
    x = self.MobileNet_ConvAdd_conv1(x)
    x = self.MobileNet_ConvAdd_fc2(x)
    return x