Training classifier does not work for more than two classes
Author: satyakesav
Pull request: no

I am trying to run a classifier on the AGN data which has four classes. I am using the following command to train and evaluate the classifier.

python examples/run_classifier.py \
  --task_name agn \
  --do_train \
  --do_eval \
  --do_lower_case \
  --data_dir $GLUE_DIR/AGN/ \
  --bert_model bert-base-uncased \
  --max_seq_length 128 \
  --train_batch_size 32 \
  --learning_rate 2e-5 \
  --num_train_epochs 2.0 \
  --output_dir /tmp/agn_output/

I have created a task named agn similar to cola, mnli and others. The model is trained properly but during evaluation it throws the following error.
'''
/pytorch/aten/src/THCUNN/ClassNLLCriterion.cu:105: void cunn_ClassNLLCriterion_updateOutput_kernel(Dtype *, Dtype *, Dtype *, long *, Dtype *, int, int, int, int, long) [with Dtype = float, Acctype = float]: block: [0,0,0], thread: [2,0,0] Assertion `t >= 0 && t < n_classes` failed.
/pytorch/aten/src/THCUNN/ClassNLLCriterion.cu:105: void cunn_ClassNLLCriterion_updateOutput_kernel(Dtype *, Dtype *, Dtype *, long *, Dtype *, int, int, int, int, long) [with Dtype = float, Acctype = float]: block: [0,0,0], thread: [3,0,0] Assertion `t >= 0 && t < n_classes` failed.
/pytorch/aten/src/THCUNN/ClassNLLCriterion.cu:105: void cunn_ClassNLLCriterion_updateOutput_kernel(Dtype *, Dtype *, Dtype *, long *, Dtype *, int, int, int, int, long) [with Dtype = float, Acctype = float]: block: [0,0,0], thread: [4,0,0] Assertion `t >= 0 && t < n_classes` failed.
/pytorch/aten/src/THCUNN/ClassNLLCriterion.cu:105: void cunn_ClassNLLCriterion_updateOutput_kernel(Dtype *, Dtype *, Dtype *, long *, Dtype *, int, int, int, int, long) [with Dtype = float, Acctype = float]: block: [0,0,0], thread: [5,0,0] Assertion `t >= 0 && t < n_classes` failed.
/pytorch/aten/src/THCUNN/ClassNLLCriterion.cu:105: void cunn_ClassNLLCriterion_updateOutput_kernel(Dtype *, Dtype *, Dtype *, long *, Dtype *, int, int, int, int, long) [with Dtype = float, Acctype = float]: block: [0,0,0], thread: [6,0,0] Assertion `t >= 0 && t < n_classes` failed.
/pytorch/aten/src/THCUNN/ClassNLLCriterion.cu:105: void cunn_ClassNLLCriterion_updateOutput_kernel(Dtype *, Dtype *, Dtype *, long *, Dtype *, int, int, int, int, long) [with Dtype = float, Acctype = float]: block: [0,0,0], thread: [7,0,0] Assertion `t >= 0 && t < n_classes` failed.
Traceback (most recent call last):
  File "examples/run_classifier.py", line 690, in <module>
    main()
  File "examples/run_classifier.py", line 663, in main
    logits = logits.detach().cpu().numpy()
RuntimeError: CUDA error: device-side assert triggered
'''

The reason for this issue is:
The model is trained with output size of 4 (since four classes), but during testing the model has output size of 2 because the BertForSequenceClassification class has default value for num_labels as 2.

So, if we change the following line in run_classifier.py
model = BertForSequenceClassification.from_pretrained(args.bert_model, state_dict=model_state_dict)
to 
model = BertForSequenceClassification.from_pretrained(args.bert_model, state_dict=model_state_dict, num_labels=num_labels), the issue will be resolved.

Please let me know If I can push the changes.

Comments:
What version of `pytorch-pretrained-BERT` are you using?

It seems to me that the change you are describing is already implemented.

https://github.com/huggingface/pytorch-pretrained-BERT/blob/0a9d7c7edb20a3e82cfbb4b72515575543784823/examples/run_classifier.py#L560

Okay. It is my bad that I did not have the latest version while debugging the issue. Thanks for pointing though. I will close the issue.