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.
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.