Type of dataset in Trainer class

sgugger
Mon, 20 Jul 2020 13:48:18 +0000

Trainer is to be used with PyTorch, so in this case the train_dataset needs to be a PyTorch dataset. TFTrainer would expect a TF dataset. The doc page is a bit unclear (types are right in the signature but too short/wrong in the enumeration). I’ll send a fix to this later today.