Well I got it “working” in that there are no errors now, but surprisingly to me, I am seeing the validation loss increase with every epoch. I’m not sure what I’m doing wrong.
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
print(f'Device used for training: {device}')
model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=1)
model.to(device)
model.train()

optim = AdamW(model.parameters(), lr=5e-5)

for epoch in range(10):
    for batch in tqdm.notebook.tqdm(train_dataloader):
        optim.zero_grad()
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        # token_type_ids = batch['token_type_ids'].to(device)
        labels = batch['labels'].to(device)
        outputs = model(input_ids=input_ids, attention_mask=attention_mask, 
                        labels=labels)
        loss = outputs[0]
        loss.backward()
        optim.step()
    validation_loss = validate(test_dataset, model)
    print(f"Validation loss in epoch {epoch}: {validation_loss}")

Here is my validate function:
def validate(test_dataset, model):
  torch.cuda.empty_cache()

  test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=128)
  total_loss = 0
  with torch.no_grad():
    for batch in tqdm.notebook.tqdm(test_dataloader):
      input_ids = batch['input_ids'].to(device)
      attention_mask = batch['attention_mask'].to(device)
      # token_type_ids = batch['token_type_ids'].to(device)
      labels = batch['labels'].to(device)
      outputs = model(input_ids=input_ids, attention_mask=attention_mask, 
                      labels=labels)
      total_loss += outputs[0].item()
  return(total_loss)

Here is some output showing the validation loss increasing:
`Device used for training: cuda`
    100%

    2871/2871 [1:49:25<00:00, 2.29s/it]

    100%

    1014/1014 [09:34<00:00, 1.76it/s]

    Validation loss in epoch 0: 70428.39735794067

    100%

    2871/2871 [2:00:01<00:00, 2.51s/it]

    100%

    1014/1014 [09:47<00:00, 1.73it/s]

    Validation loss in epoch 1: 69269.67090129852

    100%

    2871/2871 [2:01:22<00:00, 2.54s/it]

    100%

    1014/1014 [10:04<00:00, 1.68it/s]

    Validation loss in epoch 2: 70188.32639312744

    100%

    2871/2871 [1:46:28<00:00, 2.23s/it]

    100%

    1014/1014 [09:35<00:00, 1.76it/s]

    Validation loss in epoch 3: 72369.78367424011

    100%

    2871/2871 [1:47:28<00:00, 2.25s/it]

    100%

    1014/1014 [10:30<00:00, 1.61it/s]

    Validation loss in epoch 4: 73700.79190158844

    100%

    2871/2871 [1:46:24<00:00, 2.22s/it]

    100%

    1014/1014 [09:24<00:00, 1.79it/s]

    Validation loss in epoch 5: 74986.73181152344