Hi
Image (3, 256, 256). This is image passed in  batch size of 80.
I use fastai without making any changes to original model except the one above where i change the class from 1000 to one…
def se_resnext50(pretrained=False):
    pretrained = 'imagenet' if pretrained else None
    model = pt.se_resnext50_32x4d(pretrained=pretrained)

def _resnext_split(m): return   (m.layer3, m.avg_pool)    model.load_state_dict(torch.load('/tmp/.cache/torch/checkpoints/se_resnext50_32x4d-a260b3a4.pth'))
    return model

loss_func = MSELossFlat()

learn1 = Learner(data, 
                se_resnext50(),
                loss_func=loss_func ,
                
                metrics = [qk,r2_score,exp_rmspe], 
                path='.',callback_fns=[partial(SaveModelCallback,monitor='valid_loss',mode='min'),partial(ReduceLROnPlateauCallback, min_delta=1e-5, patience=3)])

learn1.callback_fns.append(partial(OverSamplingCallback1,weights=[1,1,1,1,1],bn=len(learn1.data.train_dl)))

learn1.model.last_linear=nn.Sequential(view(),nn.Linear(2048,1,bias=True))

apply_init(learn1.model.last_linear, nn.init.kaiming_uniform_)

learn1.split(_resnext_split) #this is to form param group for differential LR
learn1.model.cuda()
learn1.to_fp16()