oh, sorry, yes, it’s a copy-paste mistake.
I train the frozen layer before, the function is as followed:
def step_mlp(self,X,y):
        y_prime = self(X)
        loss = self.criterion(y_prime, y)
        
        self.optimize.zero_grad()
        self.zero_grad()
        
        loss.backward(retain_graph=1)
        self.optimize.step()

Then re-train the model by frozen the layers except the last layer as the code in question decription.
Frozen layer:
for layer in self.layers[0:-1]: 
        for param in self.model.fc_layers[layer].parameters():
            param.requires_grad = False

Then re-train:
y_prime = self.model(X)
loss = self.model.criterion(y_prime, y)
        
self.model.optimize.zero_grad()
self.model.zero_grad()
loss.backward(retain_graph=1)
self.model.optimize.step()