T5 model for rating prediction task

Beginners
1

I am trying to predict a rating for the user. But the output.logit is returning the shape of (4,1,32128) but I want to have (4,1,1). What should I do to solve this?

for epoch in range(1):
    total_loss = 0
    for batch in train_dataloader:
        users, movies, ratings = batch

        input_texts = [f"User: {user}, Movie: {movie}" for user, movie in zip(users, movies)]
        target_values = torch.tensor(ratings, dtype=torch.float)
#   
        input_encodings = tokenizer.batch_encode_plus(
            input_texts,
            padding='longest',
            truncation=True,
            return_tensors='pt'
        )
#        
        input_ids = input_encodings['input_ids']
        attention_mask = input_encodings['attention_mask']
        
        # Forward pass
#       print(target_values.unsqueeze(1).long())
        outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=target_values.unsqueeze(1).long())
        print(outputs.logits.shape)
        logits = outputs.logits.mean(dim=-1, keepdim=True)
#         
        logit=logits.reshape(-1)
        

        
        # Compute MSE loss
        loss = torch.nn.functional.mse_loss(logits, target_values)
        print(loss)
        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

    # Print average loss for the epoch
    average_loss = total_loss / len(train_dataloader)
    print(f"Epoch {epoch+1}/{epochs} - Average Loss: {average_loss}")