In the below code snippet, when I try and iterate through model.parameters() in order to obtain the param.grad data, I get a

AttributeError: ‘NoneType’ object has no attribute ‘data’

signifying that the backward pass, which is done via the loss.backward() did not store the gradient. This code worked for BERT and Electra, so not sure why it’s not working with XLNet. Any help would be appreciated.
def xln(input_ids, labels, epochs = 2):

    learning_rate = 0.02
    model = XLNetForTokenClassification.from_pretrained('xlnet-base-cased', num_labels = 2,
                                                output_hidden_states=True,
                                                output_attentions=True,
                                                )
    
    
    model.train()
    
    
    #zero's our gradient
    model.zero_grad()
    
    for epoch in range(epochs):
        
        for index, sequences in enumerate(input_ids):

            outputs = model(torch.tensor(sequences).unsqueeze(0),
                            token_type_ids=None, 
                            attention_mask=None,
                            labels = torch.tensor(labels[index]).unsqueeze(0)
                           )

            #""" this loss variable will point to the entire model """
            loss = outputs.loss
            """compute gradient. We should now have grad.data in model.parameters()"""
            loss.backward()
            
            for param in model.parameters():
                print('param',param.grad)
                param.data -= learning_rate * param.grad.data
            
            #zero gradient 
            model.zero_grad()
    
    return model
 
model = xln(x_inputs,  y_labels)