I think having a new TrainingArguments called length_field_name that defaults to "length" would be the best. If you want to make a PR, by all mean go ahead! The only thing to add in your _get_train_sampler is some test of whether the dataset comes from the datasets library (like is_datasets_available() and isinstance(self.train_dataset, datasets.Dataset))