#2d
import torch
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print("Device:", device)

import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F 
import torch.utils.data as DataLoader
from scipy.stats import pearsonr
import numpy as np
import torch.utils.data as data_utils
torch.set_printoptions(precision=10)

import Meter as Meter
import EarlyStopping as EarlyStopping



from transformers import AutoTokenizer, BertForSequenceClassification, BertTokenizerFast
model_name=MODEL_NAME
max_length=1024


from transformers import AutoTokenizer, AutoModelForTokenClassification, BertTokenizerFast, EvalPrediction, BertForSequenceClassification
from torch.utils.data import Dataset

class MyDataset(Dataset):

    def __init__(self, df, tokenizer_name=model_name, max_length=1024):
        
        #self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, do_lower_case=False)
        self.tokenizer = BertTokenizerFast.from_pretrained(tokenizer_name, do_lower_case=False)
        self.seqs, self.labels = self.load_dataset(df)
        self.max_length = max_length

    def __len__(self):
        return len(self.labels)
    def load_dataset(self,df):
        seq = list(df['sequence'])
        label = list(df['target_scaled'])
        assert len(seq) == len(label)
        return seq, label
        
    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()

        seq = " ".join("".join(self.seqs[idx].split()))
        seq = re.sub(r"[UZOB]", "X", seq)

        #seq_ids = self.tokenizer(seq, truncation=True, padding='max_length', max_length=self.max_length, return_tensors='pt').to(device)
        seq_ids = self.tokenizer(seq, truncation=True, padding='max_length', max_length=self.max_length)
        sample = {key: torch.tensor(val) for key, val in seq_ids.items()}
        sample['labels'] = torch.tensor(self.labels[idx])

        return sample

train_seqs_encodings_dataset=MyDataset(df=train_dataset_clean)
valid_seqs_encodings_dataset=MyDataset(df=valid_dataset_clean)

train_loader = DataLoader.DataLoader(
            train_seqs_encodings_dataset,
            batch_size=64,
            shuffle=True
        )
valid_loader = DataLoader.DataLoader(
            valid_seqs_encodings_dataset,
            batch_size=64,
            shuffle=True
        )


torch.manual_seed(100)
learning_rate=0.001

#Initialize network
#model=AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=1)
model=BertForSequenceClassification.from_pretrained(model_name, num_labels=1)
#Loss and optimizer
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate) #to check 
model.train()


metric_name='r2'
num_epochs=5
stopper = EarlyStopping.EarlyStopping(mode='higher', patience=20)

for epoch in range(num_epochs):
    loss_train = 0
    print(f"Epoch: {epoch+1}/{num_epochs}")
    model.train()

    train_meter = Meter.Meter()

    for batch_idx, batch in enumerate(train_loader):
        #print (batch)
        b_input_ids = batch['input_ids'].to(device=device)
        b_token_type_ids = batch['token_type_ids'].to(device=device)
        b_input_mask = batch['attention_mask'].to(device=device)
        b_labels = batch['labels'].to(device=device)

        #b_input_ids=torch.squeeze(b_input_ids)
        
        print (b_input_ids.size())
        #forward
        #predictions = model(data)#predictions=model(data.float())
        loss, predictions = model(b_input_ids, 
                             token_type_ids=b_token_type_ids, 
                             attention_mask=b_input_mask, 
                             labels=b_labels)
        
        loss= (criterion(predictions,targets)).mean()#.to(device)

        optimizer.zero_grad()

        #backward
        loss.backward()
        
        #gradient descent or adam step
        optimizer.step()
        
        loss_train += loss.item()
        train_meter.update(predictions, targets)

        score = np.mean(train_meter.compute_metric(metric_name))
    with torch.no_grad():
        val_score =  run_an_eval_epoch(valid_loader, model, metric_name)
        early_stop = stopper.step(val_score, model)

    total_score = np.mean(train_meter.compute_metric(metric_name))
    print(f'training {metric_name}: {total_score:.4f}, validation: {val_score:.4f} , best validation: {stopper.best_score:.4f}. ')
    
    if early_stop:
        break


This is the code, and the output:
Epoch: 1/5
torch.Size([64, 1, 1024])

/usr/local/anaconda/lib/python3.6/site-packages/ipykernel_launcher.py:29: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-36-0577dee49ad5> in <module>
     26                              token_type_ids=b_token_type_ids,
     27                              attention_mask=b_input_mask,
---> 28                              labels=b_labels)
     29 
     30         loss= (criterion(predictions,targets)).mean()#.to(device)

/usr/local/anaconda/lib/python3.6/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1049         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1050                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1051             return forward_call(*input, **kwargs)
   1052         # Do not call functions when jit is used
   1053         full_backward_hooks, non_full_backward_hooks = [], []

/usr/local/anaconda/lib/python3.6/site-packages/transformers/models/bert/modeling_bert.py in forward(self, input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, labels, output_attentions, output_hidden_states, return_dict)
   1509             output_attentions=output_attentions,
   1510             output_hidden_states=output_hidden_states,
-> 1511             return_dict=return_dict,
   1512         )
   1513 

/usr/local/anaconda/lib/python3.6/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1049         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1050                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1051             return forward_call(*input, **kwargs)
   1052         # Do not call functions when jit is used
   1053         full_backward_hooks, non_full_backward_hooks = [], []

/usr/local/anaconda/lib/python3.6/site-packages/transformers/models/bert/modeling_bert.py in forward(self, input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, past_key_values, use_cache, output_attentions, output_hidden_states, return_dict)
    923         elif input_ids is not None:
    924             input_shape = input_ids.size()
--> 925             batch_size, seq_length = input_shape
    926         elif inputs_embeds is not None:
    927             input_shape = inputs_embeds.size()[:-1]

ValueError: too many values to unpack (expected 2)