ValueError batch-size mismatch when redefining classifier layer on BertForSequenceClassification
Hi,

I am currently using BertForSequenceClassification for my project, to show some results regarding transfer performance on the GLUE Benchmark.

I wan't to do two things.

1. Add a seperate nn.Linear() head on top of the already fine-tuned BertForSequenceClassification model and train the entire model:
Input -> BERT_BASE_MODEL -> CLASSIFIER -> nn.Linear

2. Remove (or reinitialize) classifier head with random weights with |out_features| = |labels of the new task| of the classifier.layer of BertForSequenceClassification model and retrain it on the new task.

I have problem with 2. If I try to execute my script, I always get the following error.
I tried various things, but couldn't get the code work. I also searched for similar posts, but couldn't find one.

 >> source t_ft.sh 
Selected cpu as device.
b'Skipping line 24810: expected 12 fields, saw 13\nSkipping line 33961: expected 12 fields, saw 13\n'
b'Skipping line 75911: expected 12 fields, saw 13\nSkipping line 100114: expected 12 fields, saw 13\n'
b'Skipping line 150638: expected 12 fields, saw 13\nSkipping line 158834: expected 12 fields, saw 13\nSkipping line 173104: expected 12 fields, saw 13\nSkipping line 178252: expected 12 fields, saw 13\n'
b'Skipping line 221951: expected 12 fields, saw 13\n'
b'Skipping line 286845: expected 12 fields, saw 13\nSkipping line 314110: expected 12 fields, saw 13\n'
Processing 1000 / 391120 Samples
Processing 2000 / 391120 Samples
Processing 3000 / 391120 Samples
Processing 1000 / 9714 Samples
Processing 2000 / 9714 Samples
Processing 3000 / 9714 Samples
add_head: no
remove_head: yes

>>======== Epoch 1 / 2 ========
Training...
Traceback (most recent call last):
  File "bert_pipeline.py", line 1117, in <module>
    main()
  File "bert_pipeline.py", line 976, in main
    outputs = model(b_input_ids, token_type_ids=None, attention_mask=b_input_mask, 
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/transformers/models/bert/modeling_bert.py", line 1513, in forward
    loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/loss.py", line 1047, in forward
    return F.cross_entropy(input, target, weight=self.weight,
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/functional.py", line 2693, in cross_entropy
    return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/functional.py", line 2384, in nll_loss
    raise ValueError(
ValueError: Expected input batch_size (24) to match target batch_size (16).


I guess the error is in line 891-892 
```
if (remove_head == 'yes'):
            model.classifier = nn.Linear(in_features=model.classifier.in_features, out_features=num_labels)
```
Full Code:
```
import numpy as np_
import pandas as pd
import torch
import torch.nn as nn
from torch.nn import CrossEntropyLoss, MSELoss
import random
import time
import datetime
import argparse
import copy
import json
import csv

from transformers import BertConfig, BertTokenizer, BertForSequenceClassification, get_linear_schedule_with_warmup, AdamW
from transformers.data import metrics
from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler
from sklearn.metrics import f1_score, matthews_corrcoef
from scipy.stats import pearsonr, spearmanr


import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)

# Methods for Transfer-Learning
def freeze_base_weights(model):
    pass

# Types of some new BERT architectures

class BertWithAdditionalHead(nn.Module):
    
    def __init__(self,base_model, num_labels):
        super(BertWithAdditionalHead,self).__init__()
        self.num_labels = num_labels
        self.base_model = base_model
        self.activation = nn.GELU()
        self.fc1 = nn.Linear(self.base_model.num_labels, self.num_labels)
    
    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        labels=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
            Labels for computing the sequence classification/regression loss.
            Indices should be in :obj:`[0, ..., config.num_labels - 1]`.
            If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
            If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
        """
        return_dict = return_dict if return_dict is not None else self.base_model.config.use_return_dict

        outputs = self.base_model.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        pooled_output = outputs[1]

        pooled_output = self.base_model.dropout(pooled_output)
        outputs = self.base_model.classifier(pooled_output)
        outputs = self.activation(outputs)
        logits = self.fc1(outputs)

        loss = None
        if labels is not None:
            if self.num_labels == 1:
                #  We are doing regression
                loss_fct = MSELoss()
                loss = loss_fct(logits.view(-1), labels.view(-1))
            else:
                loss_fct = CrossEntropyLoss()
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))

        if not return_dict:
            output = (logits, outputs[2:])
            return (loss, output) if loss is not None else output

        return ((loss,logits))

# Processors     
class ColaProcessor:
    def __init__(self, data_dir):
        self.data_dir = data_dir
        self.labels=["0","1"]
        self.train_label_index=[1]
        self.dev_label_index=[1]
        
        self.train_sentence_index=[3]
        self.dev_sentence_index=[3]
        self.test_sentence_index=[1]
        
    def get_train_data(self):
        data = pd.read_csv(self.data_dir + "train.tsv",
                           delimiter="\t",
                           error_bad_lines=False,
                           header=None,
                           encoding='utf8',
                           dtype=str)
        #Remove NaN values
        data = data.dropna(subset=(self.train_sentence_index + self.train_label_index))
        train_data = data.iloc[:,self.train_sentence_index].copy()
        train_labels = data.iloc[:,self.train_label_index].copy()
        return((train_data, train_labels))

    def get_dev_data(self):
        data = pd.read_csv(self.data_dir + "dev.tsv",
                           delimiter="\t",
                           error_bad_lines=False,
                           encoding='utf8',
                           header=None,
                           dtype=str)
        data = data.dropna(subset=(self.dev_sentence_index + self.dev_label_index))
        dev_data = data.iloc[:,self.dev_sentence_index].copy()
        dev_labels = data.iloc[:,self.dev_label_index].copy()
        return((dev_data, dev_labels))
        
    def get_test_data(self):
        data = pd.read_csv(self.data_dir + "test.tsv",
                           delimiter="\t",
                           encoding='utf8',
                           error_bad_lines=False,
                           dtype=str)
        data = data.dropna(subset=self.test_sentence_index)
        test_data  = data.iloc[:,self.test_sentence_index].copy()
        return(test_data)
        
    def get_label_list(self):
        return(self.labels)
        
    def get_index(self):
        return((self.train_sentence_index, self.train_label_index),
               (self.dev_sentence_index, self.dev_label_index),
               (self.test_sentence_index))


# TODO
class MRPCProcessor:
    def __init__(self, data_dir):
        self.data_dir = data_dir
        self.labels=["0","1"]
        
        self.train_label_index=[0]
        self.dev_label_index=[0]
        
        self.train_sentence_index=[3,4]
        self.dev_sentence_index=[3,4]
        self.test_sentence_index=[3,4]
        
    def get_train_data(self):
        data = pd.read_csv(self.data_dir + "train.tsv",
                           delimiter="\t", 
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           encoding='utf8',
                           error_bad_lines=False)
        #Remove NaN values
        data = data.dropna(subset=(self.train_sentence_index + self.train_label_index))        
        train_data = data.iloc[:,self.train_sentence_index].copy()
        train_labels = data.iloc[:,self.train_label_index].copy()
        return((train_data, train_labels))

    def get_dev_data(self):
        data = pd.read_csv(self.data_dir + "dev.tsv",
                           delimiter="\t", 
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           encoding='utf8',
                           error_bad_lines=False)
        data = data.dropna(subset=(self.dev_sentence_index + self.dev_label_index))
        dev_data = data.iloc[:,self.dev_sentence_index].copy()
        dev_labels = data.iloc[:,self.dev_label_index].copy()
        return((dev_data, dev_labels))
        
    def get_test_data(self):
        data = pd.read_csv(self.data_dir + "test.tsv", 
                           delimiter="\t",
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           encoding='utf8',
                           error_bad_lines=False)
        data = data.dropna(subset=self.test_sentence_index)
        test_data  = data.iloc[:,self.test_sentence_index].copy()
        return(test_data)
        
    def get_label_list(self):
        return(self.labels)
        
    def get_index(self):
        return((self.train_sentence_index, self.train_label_index),
               (self.dev_sentence_index, self.dev_label_index),
               (self.test_sentence_index))
        
class MNLIMatchedProcessor:
    def __init__(self, data_dir):
        self.data_dir = data_dir
        self.labels=["contradiction", "entailment", "neutral"]
        
        self.train_label_index=[11]
        self.dev_label_index=[15]
        
        self.train_sentence_index=[8,9]
        self.dev_sentence_index=[8,9]
        self.test_sentence_index=[8,9]
        
    def get_train_data(self):
        data = pd.read_csv(self.data_dir + "train.tsv",
                           delimiter="\t", 
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           encoding='utf8',
                           error_bad_lines=False)
        #Remove NaN values
        data = data.dropna(subset=(self.train_sentence_index + self.train_label_index))        
        train_data = data.iloc[:,self.train_sentence_index].copy()
        train_labels = data.iloc[:,self.train_label_index].copy()
        return((train_data, train_labels))

    def get_dev_data(self):
        data = pd.read_csv(self.data_dir + "dev_matched.tsv",
                           delimiter="\t", 
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           encoding='utf8',
                           error_bad_lines=False)
        data = data.dropna(subset=(self.dev_sentence_index + self.dev_label_index))
        dev_data = data.iloc[:,self.dev_sentence_index].copy()
        dev_labels = data.iloc[:,self.dev_label_index].copy()
        return((dev_data, dev_labels))
        
    def get_test_data(self):
        data = pd.read_csv(self.data_dir + "test_matched.tsv", 
                           delimiter="\t",
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           encoding='utf8',
                           error_bad_lines=False)
        data = data.dropna(subset=self.test_sentence_index)
        test_data  = data.iloc[:,self.test_sentence_index]
        return(test_data)
        
    def get_label_list(self):
        return(self.labels)
        
    def get_index(self):
        return((self.train_sentence_index, self.train_label_index),
               (self.dev_sentence_index, self.dev_label_index),
               (self.test_sentence_index))

class QNLIProcessor:
    def __init__(self, data_dir):
        self.data_dir = data_dir
        self.labels=["entailment", "not_entailment"]
        
        self.train_label_index=[3]
        self.dev_label_index=[3]
        
        self.train_sentence_index=[1,2]
        self.dev_sentence_index=[1,2]
        self.test_sentence_index=[1,2]
        
    def get_train_data(self):
        data = pd.read_csv(self.data_dir + "train.tsv",
                           delimiter="\t", 
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           encoding='utf8',
                           error_bad_lines=False)
        #Remove NaN values
        data = data.dropna(subset=(self.train_sentence_index + self.train_label_index))        
        train_data = data.iloc[:,self.train_sentence_index].copy()
        train_labels = data.iloc[:,self.train_label_index].copy()
        return((train_data, train_labels))

    def get_dev_data(self):
        data = pd.read_csv(self.data_dir + "dev.tsv",
                           delimiter="\t", 
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           encoding='utf8',
                           error_bad_lines=False)
        data = data.dropna(subset=(self.dev_sentence_index + self.dev_label_index))
        dev_data = data.iloc[:,self.dev_sentence_index].copy()
        dev_labels = data.iloc[:,self.dev_label_index].copy()
        return((dev_data, dev_labels))
        
    def get_test_data(self):
        data = pd.read_csv(self.data_dir + "test.tsv", 
                           delimiter="\t",
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           encoding='utf8',
                           error_bad_lines=False)
        data = data.dropna(subset=self.test_sentence_index)
        test_data  = data.iloc[:,self.test_sentence_index]
        return(test_data)
        
    def get_label_list(self):
        return(self.labels)
        
    def get_index(self):
        return((self.train_sentence_index, self.train_label_index),
               (self.dev_sentence_index, self.dev_label_index),
               (self.test_sentence_index))
        
class QQPProcessor:
    def __init__(self, data_dir):
        self.data_dir = data_dir
        self.labels=["0", "1"]
        
        self.train_label_index=[5]
        self.dev_label_index=[5]
        
        self.train_sentence_index=[3,4]
        self.dev_sentence_index=[3,4]
        self.test_sentence_index=[1,2]
        
    def get_train_data(self):
        data = pd.read_csv(self.data_dir + "train.tsv",
                           delimiter="\t", 
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           encoding='utf8',
                           error_bad_lines=False)
        #Remove NaN values
        data = data.dropna(subset=(self.train_sentence_index + self.train_label_index))        
        train_data = data.iloc[:,self.train_sentence_index].copy()
        train_labels = data.iloc[:,self.train_label_index].copy()
        return((train_data, train_labels))

    def get_dev_data(self):
        data = pd.read_csv(self.data_dir + "dev.tsv",
                           delimiter="\t", 
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           encoding='utf8',
                           error_bad_lines=False)
        data = data.dropna(subset=(self.dev_sentence_index + self.dev_label_index))
        dev_data = data.iloc[:,self.dev_sentence_index].copy()
        dev_labels = data.iloc[:,self.dev_label_index].copy()
        return((dev_data, dev_labels))
        
    def get_test_data(self):
        data = pd.read_csv(self.data_dir + "test.tsv", 
                           delimiter="\t",
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           encoding='utf8',
                           error_bad_lines=False)
        data = data.dropna(subset=self.test_sentence_index)
        test_data  = data.iloc[:,self.test_sentence_index]
        return(test_data)
        
    def get_label_list(self):
        return(self.labels)
        
    def get_index(self):
        return((self.train_sentence_index, self.train_label_index),
               (self.dev_sentence_index, self.dev_label_index),
               (self.test_sentence_index))
        
class RTEProcessor:
    def __init__(self, data_dir):
        self.data_dir = data_dir
        self.labels=["entailment", "not_entailment"]
        
        self.train_label_index=[3]
        self.dev_label_index=[3]
        
        self.train_sentence_index=[1,2]
        self.dev_sentence_index=[1,2]
        self.test_sentence_index=[1,2]
        
    def get_train_data(self):
        data = pd.read_csv(self.data_dir + "train.tsv",
                           delimiter="\t", 
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           encoding='utf8',
                           error_bad_lines=False)
        #Remove NaN values
        data = data.dropna(subset=(self.train_sentence_index + self.train_label_index))        
        train_data = data.iloc[:,self.train_sentence_index].copy()
        train_labels = data.iloc[:,self.train_label_index].copy()
        return((train_data, train_labels))

    def get_dev_data(self):
        data = pd.read_csv(self.data_dir + "dev.tsv",
                           delimiter="\t", 
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           encoding='utf8',
                           error_bad_lines=False)
        data = data.dropna(subset=(self.dev_sentence_index + self.dev_label_index))
        dev_data = data.iloc[:,self.dev_sentence_index].copy()
        dev_labels = data.iloc[:,self.dev_label_index].copy()
        return((dev_data, dev_labels))
        
    def get_test_data(self):
        data = pd.read_csv(self.data_dir + "test.tsv", 
                           delimiter="\t",
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           encoding='utf8',
                           error_bad_lines=False)
        data = data.dropna(subset=self.test_sentence_index)
        test_data  = data.iloc[:,self.test_sentence_index]
        return(test_data)
        
    def get_label_list(self):
        return(self.labels)
        
    def get_index(self):
        return((self.train_sentence_index, self.train_label_index),
               (self.dev_sentence_index, self.dev_label_index),
               (self.test_sentence_index))
        
class SST2Processor:
    def __init__(self, data_dir):
        self.data_dir = data_dir
        self.labels=["0", "1"]
        
        self.train_label_index=[1]
        self.dev_label_index=[1]
        
        self.train_sentence_index=[0]
        self.dev_sentence_index=[0]
        self.test_sentence_index=[1]
        
    def get_train_data(self):
        data = pd.read_csv(self.data_dir + "train.tsv",
                           delimiter="\t", 
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           encoding='utf8',
                           error_bad_lines=False)
        #Remove NaN values
        data = data.dropna(subset=(self.train_sentence_index + self.train_label_index))        
        train_data = data.iloc[:,self.train_sentence_index].copy()
        train_labels = data.iloc[:,self.train_label_index].copy()
        return((train_data, train_labels))

    def get_dev_data(self):
        data = pd.read_csv(self.data_dir + "dev.tsv",
                           delimiter="\t", 
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           encoding='utf8',
                           error_bad_lines=False)
        data = data.dropna(subset=(self.dev_sentence_index + self.dev_label_index))
        dev_data = data.iloc[:,self.dev_sentence_index].copy()
        dev_labels = data.iloc[:,self.dev_label_index].copy()
        return((dev_data, dev_labels))
        
    def get_test_data(self):
        data = pd.read_csv(self.data_dir + "test.tsv", 
                           delimiter="\t",
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           encoding='utf8',
                           error_bad_lines=False)
        data = data.dropna(subset=self.test_sentence_index)
        test_data  = data.iloc[:,self.test_sentence_index]
        return(test_data)
        
    def get_label_list(self):
        return(self.labels)
        
    def get_index(self):
        return((self.train_sentence_index, self.train_label_index),
               (self.dev_sentence_index, self.dev_label_index),
               (self.test_sentence_index))
class STSBProcessor:
    def __init__(self, data_dir):
        self.data_dir = data_dir
        self.labels=[]
        
        self.train_label_index=[9]
        self.dev_label_index=[9]
        
        self.train_sentence_index=[7,8]
        self.dev_sentence_index=[7,8]
        self.test_sentence_index=[7,8]
        
    def get_train_data(self):
        data = pd.read_csv(self.data_dir + "train.tsv",
                           delimiter="\t", 
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           encoding='utf8',
                           error_bad_lines=False)
        #Remove NaN values
        data = data.dropna(subset=(self.train_sentence_index + self.train_label_index))        
        train_data = data.iloc[:,self.train_sentence_index].copy()
        train_labels = data.iloc[:,self.train_label_index].copy()
        return((train_data, train_labels))

    def get_dev_data(self):
        data = pd.read_csv(self.data_dir + "dev.tsv",
                           delimiter="\t", 
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           encoding='utf8',
                           error_bad_lines=False)
        data = data.dropna(subset=(self.dev_sentence_index + self.dev_label_index))
        dev_data = data.iloc[:,self.dev_sentence_index].copy()
        dev_labels = data.iloc[:,self.dev_label_index].copy()
        return((dev_data, dev_labels))
        
    def get_test_data(self):
        data = pd.read_csv(self.data_dir + "test.tsv", 
                           delimiter="\t",
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           quoting=csv.QUOTE_NONE,
                           encoding='utf8',
                           error_bad_lines=False)
        data = data.dropna(subset=self.test_sentence_index)
        test_data  = data.iloc[:,self.test_sentence_index]
        return(test_data)
        
    def get_label_list(self):
        return(self.labels)
        
    def get_index(self):
        return((self.train_sentence_index, self.train_label_index),
               (self.dev_sentence_index, self.dev_label_index),
               (self.test_sentence_index))
        
class WNLIProcessor:
    def __init__(self, data_dir):
        self.data_dir = data_dir
        self.labels=["0", "1"]
        
        self.train_label_index=[3]
        self.dev_label_index=[3]
        
        self.train_sentence_index=[1,2]
        self.dev_sentence_index=[1,2]
        self.test_sentence_index=[1,2]
        
    def get_train_data(self):
        data = pd.read_csv(self.data_dir + "train.tsv",
                           delimiter="\t", 
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           encoding='utf8',
                           error_bad_lines=False)
        #Remove NaN values
        data = data.dropna(subset=(self.train_sentence_index + self.train_label_index))
        train_data = data.iloc[:,self.train_sentence_index].copy()
        train_labels = data.iloc[:,self.train_label_index].copy()
        return((train_data, train_labels))

    def get_dev_data(self):
        data = pd.read_csv(self.data_dir + "dev.tsv",
                           delimiter="\t", 
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           encoding='utf8',
                           error_bad_lines=False)
        data = data.dropna(subset=(self.dev_sentence_index + self.dev_label_index))
        dev_data = data.iloc[:,self.dev_sentence_index].copy()
        dev_labels = data.iloc[:,self.dev_label_index].copy()
        return((dev_data, dev_labels))
        
    def get_test_data(self):
        data = pd.read_csv(self.data_dir + "test.tsv", 
                           delimiter="\t",
                           dtype=str,
                           header=None,
                           skiprows=[0],
                           quoting=csv.QUOTE_NONE,
                           encoding='utf8',
                           error_bad_lines=False)
        data = data.dropna(subset=self.test_sentence_index)
        test_data  = data.iloc[:,self.test_sentence_index]
        return(test_data)
        
    def get_label_list(self):
        return(self.labels)
        
    def get_index(self):
        return((self.train_sentence_index, self.train_label_index),
               (self.dev_sentence_index, self.dev_label_index),
               (self.test_sentence_index))
# Metrics
class Metrics():
    def __init__(self, is_regression):
        self.is_regression = is_regression
    
    def get_dict(self):
        if self.is_regression:
            return({"pearson_corr": None, "spearman_corr": None})
        else:
            return({"accuracy": None, "f1_score": None, "mcc": None})
    
    def calculate_metrics(self, preds, labels):
        eval_dict = {}
        if self.is_regression:
            eval_dict["pearson_corr"] = pearsonr(preds, labels)[0]
            eval_dict["spearman_corr"] = spearmanr(preds, labels)[0] 
        else:
            eval_dict["accuracy"] = metrics.simple_accuracy(preds, labels)
            eval_dict["f1_score"] = f1_score(y_true=labels, y_pred=preds)
            eval_dict["mcc"] = matthews_corrcoef(labels, preds)
        return(eval_dict)
        
def log_eval(epoch_i, avg_train_loss, eval_loss, eval_dict, output_dir, t_train, t_val):
    eval_file = open((output_dir + "/eval_result_transfer.txt"), "a")
    eval_file.writelines("epoch : {} \n".format(epoch_i))
    eval_file.writelines("train_loss : {} \n".format(avg_train_loss))
    eval_file.writelines("train_time : {} \n".format(t_train))
    eval_file.writelines("eval_loss : {} \n".format(eval_loss))
    eval_file.writelines("eval_time : {} \n".format(t_val))
    for key in eval_dict:
        eval_file.writelines(key + ": {} \n".format(eval_dict[key]))
    eval_file.writelines("\n")
    eval_file.close()
    with open(output_dir + "/epoch_{}.json".format(epoch_i), "w") as f:
        json.dump(eval_dict, f)
    f.close()
    

def preprocesser(dataset, sentence_idx, tokenizer, max_seq_len):
    input_ids = []
    attention_mask = []
    token_type_ids = []
    
    if len(sentence_idx) == 1:
        for j, sentence in enumerate(dataset.iloc()):
            sentence1 = sentence[sentence_idx[0]]
            
            if pd.isnull(sentence1):
                continue
            
            if j%1000==0 and not (j == 0):
                print("Processing {} / {} Samples".format(j, len(dataset)))
            # Tokenize the sentence
            sentence1 = tokenizer.tokenize(sentence1)
            
            # Convert tokens to ids
            sentence1 = tokenizer.convert_tokens_to_ids(sentence1)
            
            # Additional preprocessing
            ## Padding to a max_seq_len
            ## or Truncating to max_seq_len
                ## computing input_ids, attention mask and token_type_ids

            tokenized_dict = tokenizer.prepare_for_model(
                    ids=sentence1, pair_ids=None, add_special_tokens=True
                    ,padding='max_length', truncation='longest_first', 
                    max_length=max_seq_len, return_tensors='np', return_token_type_ids=True
                    , return_attention_mask=True)
                
            input_ids.append(tokenized_dict['input_ids'])
            attention_mask.append(tokenized_dict['attention_mask'])
            token_type_ids.append(tokenized_dict['token_type_ids'])
        
    if len(sentence_idx) == 2:
        for j,sentence in enumerate(dataset.iloc()):
            sentence1 = sentence[sentence_idx[0]]
            sentence2 = sentence[sentence_idx[1]]
            
            if pd.isnull(sentence1) or pd.isnull(sentence2):
                continue
            
            if j%1000==0 and not (j == 0):
                print("Processing {} / {} Samples".format(j, len(dataset)))
                
            if j==3000:
                break
            
            # Tokenize the sentence
            sentence1 = tokenizer.tokenize(sentence1)
            sentence2 = tokenizer.tokenize(sentence2)
            
            # Convert tokens to ids
            sentence1 = tokenizer.convert_tokens_to_ids(sentence1)
            sentence2 = tokenizer.convert_tokens_to_ids(sentence2)
            
            # Additional preprocessing
            ## Padding to a max_seq_len
            ## or Truncating to max_seq_len
            ## computing input_ids, attention mask and token_type_ids
            tokenized_dict = tokenizer.prepare_for_model(ids=sentence1, pair_ids=sentence2,
                    add_special_tokens=True,padding='max_length', truncation='longest_first',
                    max_length=max_seq_len, return_tensors='np',return_token_type_ids=True ,
                    return_attention_mask=True)
                
            input_ids.append(tokenized_dict['input_ids'])
            attention_mask.append(tokenized_dict['attention_mask'])
            token_type_ids.append(tokenized_dict['token_type_ids'])

    # Converting to pytorch tensors
    input_ids = torch.tensor(input_ids)
    attention_mask = torch.tensor(attention_mask)
    token_type_ids = torch.tensor(token_type_ids)

    return (input_ids, attention_mask, token_type_ids)

def format_time(elapsed):
    '''
    Takes a time in seconds and returns a string hh:mm:ss
    '''
    # Round to the nearest second.
    elapsed_rounded = int(round((elapsed)))
    
    # Format as hh:mm:ss
    return str(datetime.timedelta(seconds=elapsed_rounded))

def main():
    # Arguments
    parser = argparse.ArgumentParser(description='A BERT pipeline with transformers library')
    parser.add_argument('-t_n', '--task_name', help='Name of the task', default=None, type=str)
    parser.add_argument('-d_t', '--do_train', help='Whether model needs to be trained yes/no', default='no', type=str)
    parser.add_argument('-d_e', '--do_eval', help='Whether you want to evaluate on dev set', default='no', type=str)
    parser.add_argument('-d_p', '--do_predict', help='Whether you want to do predictions on test set yes/no', default='no', type=str)
    parser.add_argument('-a_h', '--add_head', help='Whether you want to add a new head on the given BERT model yes/no', default='no', type=str)
    parser.add_argument('-r_h', '--remove_head', help='Whether you want to remove head and instantiante new head with random weights yes/no', 
                        default='no', type=str)
    parser.add_argument('-f_b', '--freeze_base', help="Whether you only want to train the classification layer yes/no", default='no', type=str)
    parser.add_argument('-i_r', '--is_regression', help='Whether the given task is a regression task yes/no', default='no', type=str)
    parser.add_argument('-g_s', '--global_seed', help='Define seed for reproducability purpose', default=0, type=int)
    parser.add_argument('-d_d', '--data_dir', help='Directory, where the dataset can be found', default=None, type=str)
    parser.add_argument('-v_f', '--vocab_file', help='Path of BERT vocabulary file', default=None, type=str)
    parser.add_argument('-s_t', '--source_task', help='Optional argument, to store the name of the model', default='',type=str)
    parser.add_argument('-b_c_f', '--bert_config_file', help='Directory, where the configuration file can be found', default=None, type=str)
    parser.add_argument('-p_m', '--pretrained_model', help='Path of the Pretrained model (.bin /.pth)', default=None, type=str)
    parser.add_argument('-m_s_l', '--max_seq_len', help='Maximum length boundary for all sequences', default=128, type=int)
    parser.add_argument('-t_b_s', '--train_batch_size', help='Batch size for Training', default=32, type=int)
    parser.add_argument('-e_b_s', '--eval_batch_size', help='Batch size for Evaluation', default=16, type=int)
    parser.add_argument('-l_r', '--learning_rate', help='Learning rate', default=3e-5, type=float)
    parser.add_argument('-n_t_e', '--num_train_epochs', help='Number of training epochs', default=1, type=int)
    parser.add_argument('-n_w_s', '--num_warmup_steps', help='Number of warmup steps', default=0, type=int)
    parser.add_argument('-o_d', '--output_dir', help='Directory for the output file', default=None, type=str)
    args = vars(parser.parse_args())
    
    # Passing arguments to variables
    task_name = args['task_name']
    do_train = args['do_train']
    do_eval = args['do_eval']
    do_predict = args['do_predict']
    add_head = args['add_head']
    remove_head = args['remove_head']
    freeze_base = args['freeze_base']
    global_seed = args['global_seed']
    data_dir = args['data_dir']
    vocab_file = args['vocab_file']
    bert_config_file = args['bert_config_file']
    pretrained_model = args['pretrained_model']
    max_seq_len = args['max_seq_len']
    train_batch_size = args['train_batch_size']
    eval_batch_size = args['eval_batch_size']
    learning_rate = args['learning_rate']
    epochs = args['num_train_epochs']
    num_warmup_steps = args['num_warmup_steps']
    output_dir = args['output_dir']
    is_regression = args['is_regression']
    source_task = args['source_task']
    # Setting seed
    random.seed(global_seed)
    np.random.seed(global_seed)
    torch.manual_seed(global_seed)
    torch.cuda.manual_seed_all(global_seed)
    
    # Setting up device
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    print("Selected {} as device.".format(device))
    
    # BertConfig
    bert_config = BertConfig.from_json_file(bert_config_file)
    
    # Processor
    processor_dict = {"cola":ColaProcessor,
                      "mrpc":MRPCProcessor,
                      "mnli":MNLIMatchedProcessor,
                      "qnli":QNLIProcessor,
                      "qqp":QQPProcessor,
                      "rte":RTEProcessor,
                      "sst2":SST2Processor,
                      "sts":STSBProcessor,
                      "wnli":WNLIProcessor}
    
    processor = processor_dict[task_name](data_dir=data_dir)
    
    # Tokenizer
    tokenizer = BertTokenizer(vocab_file=vocab_file)
    
    # Metrics
    metric = Metrics(is_regression)
    
    # Training
    if do_train == 'yes':
        train_data, train_labels = processor.get_train_data()
        dev_data, dev_labels = processor.get_dev_data()
        
        label_list = processor.get_label_list()
        num_labels = len(label_list)
        (train_sentence_index, _), (dev_sentence_index, _), _ = processor.get_index()
        
        train_input_ids, train_attention_mask, train_token_type_ids = preprocesser(
                dataset=train_data, sentence_idx=train_sentence_index,tokenizer=tokenizer, 
                max_seq_len=max_seq_len)
        dev_input_ids, dev_attention_mask, dev_token_type_ids = preprocesser(
                dataset=dev_data, sentence_idx=dev_sentence_index
                ,tokenizer=tokenizer, max_seq_len=max_seq_len)
        
        
        # Converting labels to numeric values
        train_labels = train_labels.values.flatten('C')
        dev_labels = dev_labels.values.flatten('C')
        if is_regression == 'yes':
            # TODO
            train_labels = train_labels.astype(float)
            dev_labels = dev_labels.astype(float)
        else:
            label_map = {}
        
            for i,label in enumerate(label_list):
                label_map[label] = i
            
            for j,label in enumerate(train_labels):
                train_labels[j] = label_map[label]
            
            for k,label in enumerate(dev_labels):
                dev_labels[k] = label_map[label]
            
            train_labels=train_labels.astype(int)
            dev_labels=dev_labels.astype(int)
        
        train_labels = torch.tensor(train_labels[:3000])
        dev_labels = torch.tensor(dev_labels[:3000])
        
        # Defining DataLoader
        train_data = TensorDataset(train_input_ids, train_attention_mask, train_labels)
        train_sampler = RandomSampler(train_data)
        train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=train_batch_size)
        
        dev_data = TensorDataset(dev_input_ids, dev_attention_mask, dev_labels)
        dev_sampler = SequentialSampler(dev_data)
        dev_dataloader = DataLoader(dev_data, sampler=dev_sampler, batch_size=eval_batch_size)
        
        
        # TODO Transfer-Modeling
        # Defining model and criterion
        bert_config.num_labels = num_labels
        model = BertForSequenceClassification.from_pretrained(pretrained_model, return_dict=False)
        print('add_head: {}'.format(add_head))
        print('remove_head: {}'.format(remove_head))
        if (add_head == 'yes'):
            model = BertWithAdditionalHead(model, num_labels)
        **```
if (remove_head == 'yes'):
            model.classifier = nn.Linear(in_features=model.classifier.in_features, out_features=num_labels)
```**

        # TODO Freezing weights of base model
        
        if (freeze_base == 'yes'):
            if (add_head == 'yes'):
                for params in model.base_model.parameters():
                    print("Freezing Parameter: {}".format(params))
                    params.requires_grad = False
            else:
                for params in model.bert.parameters():
                    print("Freezing Parameter: {}".format(params))
                    params.requires_grad = False

        # Moving model to GPU if possible
        if device == torch.device("cuda"):
            model.cuda()
        
        optimizer = AdamW(model.parameters(), lr=learning_rate)
        total_steps = len(train_dataloader) * epochs
        scheduler = get_linear_schedule_with_warmup(optimizer,num_warmup_steps=num_warmup_steps,
                                num_training_steps=total_steps)
        loss_val = []
        # For each epoch...
        for epoch_i in range(0, epochs):

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

            # Perform one full pass over the training set.

            print("")
            print('======== Epoch {:} / {:} ========'.format(epoch_i + 1, epochs))
            print('Training...')

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

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

            # Put the model into training mode. Don't be mislead--the call to 
            # `train` just changes the *mode*, it doesn't *perform* the training.
            # `dropout` and `batchnorm` layers behave differently during training
            # vs. test (source: https://stackoverflow.com/questions/51433378/what-does-model-train-do-in-pytorch)
            model.train()

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

                    # Progress update every 40 batches.
                if step % 2 == 0 and not step == 0:
                    # Calculate elapsed time in minutes.
                    elapsed = format_time(time.time() - t0)

                    # Report progress.
                    print('  Batch {:>5,}  of  {:>5,}.    Elapsed: {:}.'.format(step,
                                len(train_dataloader), elapsed))

                # Unpack this training batch from our dataloader. 
                #
                # As we unpack the batch, we'll also copy each tensor to the GPU using the 
                # `to` method.
                #
                # `batch` contains three pytorch tensors:
                #   [0]: input ids 
                #   [1]: attention masks
                #   [2]: labels 
                b_input_ids = batch[0].to(device)
                b_input_mask = batch[1].to(device)
                b_labels = batch[2].to(device)

                # Always clear any previously calculated gradients before performing a
                # backward pass. PyTorch doesn't do this automatically because 
                # accumulating the gradients is "convenient while training RNNs". 
                # (source: https://stackoverflow.com/questions/48001598/why-do-we-need-to-call-zero-grad-in-pytorch)
                model.zero_grad()
                # optimizer.zero_grad()
                # Perform a forward pass (evaluate the model on this training batch).
                # This will return the loss (rather than the model output) because we
                # have provided the `labels`.
                # The documentation for this `model` function is here: 
                # https://huggingface.co/transformers/v2.2.0/model_doc/bert.html#transformers.BertForSequenceClassification
                outputs = model(b_input_ids, token_type_ids=None, attention_mask=b_input_mask, 
                    labels=b_labels)

                # The call to `model` always returns a tuple, so we need to pull the 
                # loss value out of the tuple.
                loss = outputs[0]
                
                # Display loss for every 10 steps
                print("Loss: {} in Step: {}".format(loss, step))


                if step%20==0 and not step==0:
                    break
                # Accumulate the training loss over all of the batches so that we can
                # calculate the average loss at the end. `loss` is a Tensor containing a
                # single value; the `.item()` function just returns the Python value 
                # from the tensor.
                total_loss += loss.item()

                # Perform a backward pass to calculate the gradients.
                loss.backward()

                # Clip the norm of the gradients to 1.0.
                # This is to help prevent the "exploding gradients" problem.
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

                # Update parameters and take a step using the computed gradient.
                # The optimizer dictates the "update rule"--how the parameters are
                # modified based on their gradients, the learning rate, etc.
                optimizer.step()

                # Update the learning rate.
                scheduler.step()

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

            # Store the loss value for plotting the learning curve.
            loss_val.append(avg_train_loss)
            
            t_train = format_time(time.time() - t0)
            print("")
            print("  Average training loss: {0:.2f}".format(avg_train_loss))
            print("  Training epoch took: {:}".format(t_train))
        
            # ========================================
            #               Validation
            # ========================================
            # After the completion of each training epoch, measure our performance on
            # our validation set.

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

            t0 = time.time()

            # Put the model in evaluation mode--the dropout layers behave differently
            # during evaluation.
            model.eval()

            # Tracking variables
            eval_loss, eval_accuracy = 0, 0
            nb_eval_steps, nb_eval_examples = 0, 0
            eval_dict = metric.get_dict()
            tmp_eval_dict = {}

            # Evaluate data for one epoch
            for dev_step, batch in enumerate(dev_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.
                    # This will return the logits rather than the loss because we have
                    # not provided labels.
                    # token_type_ids is the same as the "segment ids", which 
                    # differentiates sentence 1 and 2 in 2-sentence tasks.
                    # The documentation for this `model` function is here: 
                    # https://huggingface.co/transformers/v2.2.0/model_doc/bert.html#transformers.BertForSequenceClassification
                    outputs = model(b_input_ids,
                                    token_type_ids=None,
                                    attention_mask=b_input_mask,
                                    labels=b_labels)

                # Get the "logits" output by the model. The "logits" are the output
                # values prior to applying an activation function like the softmax.
                tmp_eval_loss, logits = outputs[:2]

                print("Eval Loss: {} in Step: {}".format(tmp_eval_loss, dev_step))
                # Move logits and labels to CPU
                logits = logits.detach().cpu().numpy()
                logits = logits.argmax(axis=1)
                label_ids = b_labels.to('cpu').numpy()

                # Calculate the accuracy for this batch of test sentences.
                # tmp_eval_accuracy = flat_accuracy(logits, label_ids)
                tmp_eval_dict = metric.calculate_metrics(preds=logits, labels=label_ids)

                # Accumulate the total accuracy.
                eval_loss += tmp_eval_loss
                for key in eval_dict:
                    if eval_dict[key] == None:
                        eval_dict = copy.deepcopy(tmp_eval_dict)
                        continue
                    else:
                        eval_dict[key] += tmp_eval_dict[key]

                # Track the number of batches
                nb_eval_steps += 1
            # logging time
                if dev_step==10:
                    break
            t_val = format_time(time.time() - t0)
            
            for key in eval_dict:
                eval_dict[key] = eval_dict[key]/nb_eval_steps
            eval_loss = eval_loss/nb_eval_steps
            log_eval(epoch_i, avg_train_loss, eval_loss, eval_dict, output_dir, t_train, t_val)
            # Report the final accuracy for this validation run.
            for key in eval_dict:
                print(key + ": {}".format(eval_dict[key]))
            print("  Validation took: {:}".format(t_val))
        print("")
        print("Training complete!")
        print("Saving the model in {} ...".format(output_dir))
        model.save_pretrained(output_dir+"{}_{}.bin".format(source_task,task_name))
    
    # TODO
    if do_eval == 'yes' and not (do_train == 'yes'):
        pass
        
   # TODO
    if do_predict == 'yes':
        pass

if __name__ == "__main__":
    main()
```
Training is based on: https://www.youtube.com/watch?v=FKlPCK1uFrc&list=PLam9sigHPGwOBuH4_4fr-XvDbe5uneaf6
Hopefully you can help me :)