Claasification Model Structure

[chiguru1]

i am new to ml, so i want to know the model clearly, so i was trying to write the model with configuration model

model = dict(
    type='ImageClassifier',
    backbone=dict(
        type='ResNet',
        depth=34,
        num_stages=4,
        out_indices=(3, ),
        style='pytorch'),
    neck=dict(type='GlobalAveragePooling'),
    head=dict(
        type='LinearClsHead',
        num_classes=12,
        in_channels=512,
        loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
        topk=(1, 5))) 

i asked chat gpt and studied a bit and wrote this as

class ImageClassification(nn.Module):
    def __init__(self, output_size=1):
        super(ImageClassification, self).__init__()
        resnet = torchvision.models.resnet50(pretrained=True)
        self.backbone = nn.Sequential(*list(resnet.children())[:-1])
        self.neck = nn.AdaptiveAvgPool2d((1,1))
        self.head = nn.Sequential(
            nn.Linear(2048, 512),
            nn.Linear(512, output_size)
        )

    def forward(self, x):
        x = self.backbone(x)
        x = self.neck(x)
        x = x.view(x.size(0), -1)
        x = self.head(x)
        return x

but it is not working at all the model is not learning
can some one help me writing this model thankyou
it will be a great help if you help me with some documentaions also is fine.