Thank you. Check the code below, I have simplified it.
#imagine we create datasets here

# initialize model
model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)
   
# This is the difference! 
if torch.cuda.device_count() > 1:
    # data parallelism
    model = torch.nn.DataParallel(model)
model.to(device)

#initialize optimizer
epochs = 3
optimizer = AdamW(model.parameters(),
      lr = 2e-5, 
      eps = 1e-8 
    )
# Total number of training steps is number of batches * number of epochs.
total_steps = len(train_dataloader) * epochs


# Create the learning rate scheduler.
scheduler = get_linear_schedule_with_warmup(optimizer, 
                                            num_warmup_steps = 0, # Default value in run_glue.py
                                            num_training_steps = total_steps)
    

# training

# Store the average loss after each epoch so we can plot them.
loss_values = []

model.zero_grad()

# For each epoch...
for epoch_i in range(0, epochs):

    # ========================================
    #               Training
    # ========================================


    # Measure how long the training epoch takes.
    t0 = time()

    # Reset the total loss for this epoch.
    total_loss = 0

    # Set our model to training mode (as opposed to evaluation mode)
    model.train()


    # For each batch of training data...
    for step, batch in enumerate(train_dataloader):


        # Unpack this training batch from our dataloader. 
        #
        b_input_ids = batch[0].to(device)
        b_input_mask = batch[1].to(device)
        b_labels = batch[2].to(device)

        # Forward pass (evaluate the model on this training batch)
        # `model` is of type: pytorch_pretrained_bert.modeling.BertForSequenceClassification
        outputs = model(b_input_ids, 
                    token_type_ids=None, 
                    attention_mask=b_input_mask, 
                    labels=b_labels)

        loss = outputs[0]

        total_loss += loss.mean().item()
#       total_loss += loss.item() # this is for single GPU

        # Perform a backward pass to calculate the gradients.
        loss.mean().backward()
#       loss.backward()  # this is for single GPU

        # Clip the norm of the gradients to 1.0.
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

        # Update parameters and take a step using the computed gradient
        optimizer.step()

        # Update the learning rate.
        scheduler.step()

        # Clear out the gradients (by default they accumulate)
        model.zero_grad()

    # Calculate the average loss over the training data.
    avg_train_loss = total_loss / len(train_dataloader)            

    loss_values.append(avg_train_loss)
    

    # ========================================
    #               Validation
    # ========================================
    # After the completion of each training epoch, measure our performance on
    # our validation set.

        print("")
        print("Running Validation...")

    t0 = time()

    # Put model in evaluation mode to evaluate loss on the validation set
    model.eval()

    # Tracking variables 
    preds, labels = [], []

    # Evaluate data for one epoch
    for batch in validation_dataloader:

        # Add batch to GPU
        batch = tuple(t.to(device) for t in batch)

        # Unpack the inputs from our dataloader
        b_input_ids, b_input_mask, b_labels = batch

        # Telling the model not to compute or store gradients, saving memory and speeding up validation
        with torch.no_grad():        
            # Forward pass, calculate logit predictions
            # token_type_ids is for the segment ids, but we only have a single sentence here.
            outputs = model(b_input_ids, token_type_ids=None, attention_mask=b_input_mask)

        logits = outputs[0]

        # Move logits and labels to CPU
        logits = logits.detach().cpu().numpy()
        label_ids = b_labels.to('cpu').numpy()

        preds.append(np.argmax(logits, axis=1))
        labels.append(label_ids)


    # Report the final accuracy for this validation run.
    preds = [item for sublist in preds for item in sublist]
    labels = [item for sublist in labels for item in sublist]    

    print_score(preds, labels) # my own function for evalution

# ========================================
#               Prediction
# ========================================

# Put model in evaluation mode
model.eval()

# Tracking variables 
predictions , true_labels = [], []

# Predict 
for batch in prediction_dataloader:
    # Add batch to GPU
    batch = tuple(t.to(device) for t in batch)

    # Unpack the inputs from our dataloader
    b_input_ids, b_input_mask, b_labels = batch

    # Telling the model not to compute or store gradients, saving memory and 
    # speeding up prediction
    with torch.no_grad():
      # Forward pass, calculate logit predictions
      outputs = model(b_input_ids, token_type_ids=None, 
                      attention_mask=b_input_mask)

    logits = outputs[0]

    # Move logits and labels to CPU
    logits = logits.detach().cpu().numpy()
    label_ids = b_labels.to('cpu').numpy()

    # Store predictions and true labels
    predictions.append(np.argmax(logits, axis=1))
    true_labels.append(label_ids)

predictions = [item for sublist in predictions for item in sublist]
true_labels = [item for sublist in true_labels for item in sublist]
print_score(predictions, true_labels, task)