Hello to everyone!
I have a pre-trained model made by myself, now I pick up some layers of the base model and use them to feed a classifier. Now I want to train the classifier and fine-tune the layers of my pre-trained net. My question is: is it correct how I would like to do here in the code below? I am not sure where there is the comment

#???

Someone can explain if, where, and why I am going wrong?
Thank you so much!
class Mybasemodel(nn.Module):
    def __init__(self, modelA, modelB):
        self.model = nn.Sequential(
            .
            .
            .
        )
    def forward(self, x1, x2):
        x = self.model(x)

        return x


class Classifier(nn.Module):
    def __init__(self, modelA, modelB):
        self.model = nn.Sequential(
            .
            .
            .
        )
    def forward(self, x1, x2):
        x = self.model(x)

        return x

basemodel = Mybasemodel()
basemodel.load_state_dict(torch.load('...'))
basemodel.train()

calssifier = Classifier()
calssifier.train()

loss_basemodel_fn = torch.nn.L1Loss()
loss_classifier_fn = torch.nn.L1Loss()

opt_basemodel = torch.optim.(basemodel.parameters())
opt_classifier = torch.optim.(calssifier.parameters())

first_layer = 0
N_layer = 3

i = 0
for param in model.basemodel():
    i += 1
    if i > N_layer:
        param.requires_grad = False


for i, data in enumerate(dataset):

    input = data['input']
    label_basemodel = data['label_basemodel']
    label_classifier = data['label_classifier']

    output_base_model = basemodel.model(input)

    feature_map = basemodel.model[first_layer:N_layer](input)
    pred = calssifier(feature_map)

    opt_basemodel.zero_grad()
    opt_classifier.zero_grad()

    loss_classifier = loss_classifier_fn(pred, label_basemodel)
    #???
    loss_basemodel = loss_basemodel_fn(output_base_model, label_classifier) + loss_classifier

    loss_basemodel.backward()
    loss_classifier.backward()

    opt_basemodel.step()
    opt_classifier.step()