Pytorch BERT input gradient
I am trying to get the input gradients from a BERT model in pytorch. How can I do that?
Suppose, y' = BertModel(x). I am trying to find $d(loss(y,y'))/dx$

One of the problems with Bert models is that your input mostly contains token IDs rather than token embeddings, which makes getting gradient difficult since the relation between token ID and token embeddings is discontinued.
To solve this issue, you can work with token embeddings.
# get your batch data: token_id, mask and labels
token_ids, mask, labels = batch
  
# get your token embeddings
token_embeds=BertModel.bert.get_input_embeddings().weight[token_ids].clone()
# track gradient of token embeddings
token_embeds.requires_grad=True
    
# get model output that contains loss value
outs = BertModel(inputs_embeds=inputs_embeds,labels=labels)
loss=outs.loss

After getting loss value, you can use torch.autograd.grad in this answer or backward function
loss.backward()
grad=token_embeds.grad

