Using XLNet for sentiment analysis - setting the correct reshape parameters

Following this link, I am trying to use my own data to do sentiment analysis. But I get this error:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<timed exec> in <module>

<ipython-input-41-5f2f35b7976e> in train_epoch(model, data_loader, optimizer, device, scheduler, n_examples)
      7 
      8     for d in data_loader:
----> 9         input_ids = d["input_ids"].reshape(4,64).to(device)
     10         attention_mask = d["attention_mask"].to(device)
     11         targets = d["targets"].to(device)

RuntimeError: shape '[4, 64]' is invalid for input of size 64

When I try to run this code

history = defaultdict(list)
best_accuracy = 0

for epoch in range(EPOCHS):
    print(f'Epoch {epoch + 1}/{EPOCHS}')
    print('-' * 10)

    train_acc, train_loss = train_epoch(
        model,
        train_data_loader,     
        optimizer, 
        device, 
        scheduler, 
        len(df_train)
    )

    print(f'Train loss {train_loss} Train accuracy {train_acc}')

    val_acc, val_loss = eval_model(
        model,
        val_data_loader, 
        device, 
        len(df_val)
    )

    print(f'Val loss {val_loss} Val accuracy {val_acc}')
    print()

    history['train_acc'].append(train_acc)
    history['train_loss'].append(train_loss)
    history['val_acc'].append(val_acc)
    history['val_loss'].append(val_loss)

I know this error has something to do with the shape of my data but I am not sure how to find the correct reshape parameters in order to make this work.

from Using XLNet for sentiment analysis - setting the correct reshape parameters