Problem with custom head with mobilenet

So, im trying to use a mobilenetv2 network from torchvision with fast ai. Im loading the model with:

import torchvision.models as models
model = models.mobilenet_v2(pretrained=True)

And now i want to change the classifier part of the model with the classic classifier from fast ai. so what im doing is:

import torch
from torch import nn
clasf = nn.Sequential(
nn.Linear(2*1280, 512),
nn.Linear(512, 37),
model.classifier = clasf   (Now my own classifier is the head of the net)

The last convolution block from the mobilenet is this one:
(18): ConvBNReLU(
(0): Conv2d(320, 1280, kernel_size=(1, 1), stride=(1, 1), bias=False)
(1): BatchNorm2d(1280, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU6(inplace=True)

So the inputs for my batchnorm1d is 2x1280 and for the linear is also 2x1280, the output is 37 for the cats/dogs classification.

So now that i have my model ready, i create my learn:
learn = Learner(data, model, loss_func = nn.CrossEntropyLoss(), metrics=accuracy)

and i get this error:
RuntimeError: non-empty 3D or 4D (batch mode) tensor expected for input

If i just delete the AdaptiveConcatPool2d() and modify the batchnorm input to 1280 (now its matching the last conv layer from mobilenet) everything works fine, but i want to understand why this is happening and what should be the input/output for the AdaptiveConcatPool2d.


This is related, but mobilenetv2 is now accessible from within fastai. See here.

Related forum post here.

Seems that this is no longer available in version 2.2.7