Use dataset.set_format in dataset collator

[cemilcengiz]

Instead of doing the tensor conversion manually, we can do something like
dataset.set_format(type='torch', columns=['input_ids', 'token_type_ids', 'attention_mask', 'label'])
as shown in datasets doc. This would also enable removing the unused columns (the columns not required by the models) conveniently.