How to build a classifier on top of trained MLM

Beginners
1

I have build an MLM using my own dataset and now I want to built a classifier using that Masked Language Model by freezing the hidden layers. So how do I remove the MLM head from the model to build a classifier. Can anyone please help me with that.

This is how I tried implementing

class BERTClass(torch.nn.Module):
    def __init__(self, bert_model):
        super(BERTClass, self).__init__()
        self.bert_model = bert_model
        self.dropout = torch.nn.Dropout(0.3)
        self.linear = torch.nn.Linear(768, 6)
    
    def forward(self, input_ids, attn_mask, token_type_ids):
        output = self.bert_model(
            input_ids=input_ids, 
            attention_mask = attn_mask,
            token_type_ids = token_type_ids
        )
        output=torch.FloatTensor(output)
        output_dropout = self.dropout(output[0])
        output = self.linear(output_dropout)
        return output
    
model = BERTClass(mlm_model)
model.to(device)

I am getting the error :
TypeError: new(): data must be a sequence (got MaskedLMOutput)

Can you please help me with this. Thank you.