Dataset map return only list instead torch tensors

I ended up using a data collator to take the lists and turn them back into tensors

1 Like