i am trying to create xlnet classification
  def __init__(self,n_classes):
    super(SentimentClassifier, self).__init__()
    self.xlnet = XLNetModel.from_pretrained(PRE_TRAINED_MODEL_NAME)
    self.drop = nn.Dropout(p=0.3)
    self.out = nn.Linear(self.xlnet.config.hidden_size,  n_classes)
  def forward(self, input_ids, 
              attention_mask):
    _, pooled_output = self.xlnet(
      input_ids=input_ids,
      attention_mask=attention_mask)
    output = self.drop(pooled_output)
    return self.out(output)
class Classification(Dataset):


    def __init__(self, texts, labels, tokenizer, max_len):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_len = max_len
    
    def __len__(self):
        return len(self.texts)
    
    def __getitem__(self, item):
        text = str(self.texts[item])
        label = self.labels[item]

        encoding = self.tokenizer.encode_plus(
        text,
        add_special_tokens=True,
        max_length=self.max_len,
        return_token_type_ids=False,
        pad_to_max_length=False,
        return_attention_mask=True,
        return_tensors='pt',
        )

        

        return {
        'review_text': text,
        'input_ids': encoding['input_ids'].flatten(),
        'attention_mask': encoding['attention_mask'].flatten(),
         'labelss': torch.tensor(label, dtype=torch.long)
        }
def train_epoch(
  model,
  data_loader,
  loss_fn,
  optimizer,
  device,
  scheduler,
  n_examples
):
  model = xlnet_model.train()
  losses = []
  correct_predictions = 0
  for d in data_loader:
    input_ids = d["input_ids"].reshape(4,512).to(device)
    print(d['input_ids'].shape)
    attention_mask = d["attention_mask"].to(device)
    labels = d["labels"].to(device)
    outputs = xlnet_model(input_ids=input_ids, attention_mask=attention_mask)
    
    _, preds = torch.max(outputs, dim=1)
    loss = loss_fn(outputs, labels)
    correct_predictions += torch.sum(preds == labels)
    losses.append(loss.item())
    loss.backward()
    nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()
    scheduler.step()
    optimizer.zero_grad()
  return correct_predictions.double() / n_examples, np.mean(loss)```