[BUG]

Describe the bug
Hey, I want to freeze part of my model in the first few epochs and unfreeze it. But it doesn't seem to work. The model doesn't seem to learn after unfreezing it, even with "torch.clear_autocast_cache() & torch.empty_cache()". Is there a valid way to make it?

To Reproduce
Bellow are the sample code of my model.py

import torch.nn as nn
import torch.nn.functional as F
from torch.nn.modules.loss import CrossEntropyLoss
from utils.mismatched_utils import MisMatchedEmbedder
from transformers import AutoModel
import torch
import time

class SeqEncoder(nn.Module):
    def __init__(self, sub_token_mode, encoder_path, device):
        super().__init__()
        self.matched_embedder = AutoModel.from_pretrained(encoder_path)
        self.hidden_size = self.matched_embedder.config.hidden_size
        self.mismatched_embedder = MisMatchedEmbedder(device, sub_token_mode)
    
    def forward(self, input_dict):
        last_hidden_states, _ = self.matched_embedder(
            input_ids=input_dict["input_ids"],
            token_type_ids=input_dict["token_type_ids"],
            attention_mask=input_dict["attention_mask"],
            return_dict=False
        )

        word_embeddings = self.mismatched_embedder.get_mismatched_embeddings(
            last_hidden_states,
            offsets=input_dict["offsets"],
            word_mask=input_dict["word_mask"])

        return word_embeddings

class GECToRModel(nn.Module):
    def __init__(self, 
        encoder_path,
        ...
        ):
        
        super().__init__()
        self.encoder = SeqEncoder(encoder_path ...)
        self._freeze_encoder = False
        ...
        
    def forward(self, input_dict):
        embeddings = self.encoder(input_dict)
        ...

        return output_dict
    
    @property
    def freeze_encoder(self):
        return self._freeze_encoder
    
    @freeze_encoder.setter
    def freeze_encoder(self, value: bool):
        for param in self.encoder.parameters():
            if value is True:
                param.requires_grad = False
            else:
                param.requires_grad = True
        self._freeze_encoder = value

Start by reproducing the model.py example and inspect GECToRModel.freeze_encoder and SeqEncoder. Check the model’s behavior before and after toggling the property, with the work complete when the encoder resumes learning after unfreezing or the valid usage is documented.