Thanks @Jung.
I am stuck at this point. I am attaching my code below.
class FC_Embeddings(nn.Module):
    def __init__(self, D_in, **kwargs):
        super(FC_Embeddings, self).__init__()
        self.D_in = D_in
        self.word_embed = nn.Embedding(30522, D_in)
        
    def forward(self, token_ids, *args, **kwargs):
        if kwargs is not None:
            fc_idxs = kwargs["fc_idxs"]
            print(fc_idxs)
        if fc_idxs is not None:
            out = self.word_embed(token_ids)
            fc_mags = torch.ones(1, 80, self.D_in)
            fc_mags[:, fc_idxs, :] *= 2
            out *= ones
        else:
            out = self.word_embed(token_ids)
        return out

%%time
import torch
import torch.nn as nn
from transformers import BertModel

# Create the BertClassfier class
class BertClassifier(nn.Module):
    """Bert Model for Classification Tasks.
    """
    def __init__(self, freeze_bert=False):
        super(BertClassifier, self).__init__()
        D_in, H, D_out = 768, 50, 5

        # Instantiate BERT model
        self.bert = BertModel.from_pretrained('bert-base-uncased')
        self.bert.base_model.embeddings.word_embeddings = FC_Embeddings(D_in)

        # Instantiate an one-layer feed-forward classifier
        self.classifier = nn.Sequential(
            nn.Linear(D_in, H),
            nn.ReLU(),
            nn.Linear(H, D_out)
        )

        # Freeze the BERT model
        if freeze_bert:
            for param in self.bert.parameters():
                param.requires_grad = False
        
    def forward(self, input_ids, attention_mask, *args, **kwargs):
        # Feed input to BERT
        outputs = self.bert(input_ids=input_ids,
                            attention_mask=attention_mask, *args, **kwargs)
        # print(outputs.shape)
        # Extract the last hidden state of the token `[CLS]` for classification task
        last_hidden_state_cls = outputs[0][:, 0, :]

        # Feed input to classifier to compute logits
        logits = self.classifier(last_hidden_state_cls)

        return logits

When I call the modified BERT I am getting this error:
token_ids = preprocessing_for_bert([X[1111]])[0]
token_mask = preprocessing_for_bert([X[1111]])[1]

b = BertClassifier()
kwargs = {"fc_idxs": [1, 2, -1]}
out = b(input_ids=token_ids, attention_mask=token_mask, **kwargs)

ERROR:
TypeError                                 Traceback (most recent call last)
<ipython-input-138-aac1926d7204> in <module>
     12 b = BertClassifier()
     13 kwargs = {"fc_idxs": [1, 2, -1]}
---> 14 out = b(input_ids=token_ids, attention_mask=token_mask, **kwargs)
     15 
     16 out

~/anaconda3/envs/aogtr/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    720             result = self._slow_forward(*input, **kwargs)
    721         else:
--> 722             result = self.forward(*input, **kwargs)
    723         for hook in itertools.chain(
    724                 _global_forward_hooks.values(),

<timed exec> in forward(self, input_ids, attention_mask, *args, **kwargs)

~/anaconda3/envs/aogtr/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    720             result = self._slow_forward(*input, **kwargs)
    721         else:
--> 722             result = self.forward(*input, **kwargs)
    723         for hook in itertools.chain(
    724                 _global_forward_hooks.values(),

TypeError: forward() got an unexpected keyword argument 'fc_idxs'

Any guidance would be very helpful!
Thanks