Fine-training inception_v3 model

Hi;

I am trying to fine-tune a pre-trained Inception v_3 model for a two class problem.

import torch
from torchvision import models
from torch.nn import nn

model = model.incepetion_v3(pretrained =True)
model.fc= nn.Linear(2048,2) ----- converting to two class problem
data = Variable (torch.rand(2,3,299,299))

outs = model(data)

outs is a tuple of dimension ([2x2] and [2x1000])

I see that the model has something called as AUX logits and I force the aux logits to have same number of neurons as the number of class as given below

model = model.incepetion_v3(pretrained =True)
model.fc= nn.Linear(2048,2) ----- converting to two class problem
model.AuxLogits=nn.Linear(768,2) – converting aux logits to 2 class

data = Variable (torch.rand(2,3,299,299))
outs = model(data)

Upon doing this I get the following error .

output = torch.mm(tensor1, tensor2)

File “/usr/local/lib/python2.7/site-packages/torch/autograd/variable.py”, line 579, in mm
return Addmm.apply(output, self, matrix, 0, 1, True)
File “/usr/local/lib/python2.7/site-packages/torch/autograd/_functions/blas.py”, line 26, in forward
matrix1, matrix2, out=output)
RuntimeError: size mismatch, m1: [26112 x 17], m2: [768 x 2] at /Users/soumith/code/builder/wheel/pytorch-src/torch/lib/TH/generic/THTensorMath.c:1293

Any ideas ?

I guess someone else too had the same issue.
Posting the answer found in another discussion here.

import torch
from torchvision import models
from torch.nn import nn

model = models.inception_v3(pretrained=True)
model.aux_logit=False

in_features = model.fc.in_features

model.fc = nn.Linear(in_features, 2)