gradient_checkpointing=True issue in TrainerArgument

[lolshuo]

I'm using the Retnet base config with the following TrainingArguments:

args = TrainingArguments( output_dir="/content/retnet-xsum", per_device_train_batch_size=1, per_device_eval_batch_size=1, evaluation_strategy="steps", eval_steps=370, logging_steps=370, num_train_epochs=10, weight_decay=0.01, warmup_steps=10, lr_scheduler_type="cosine", learning_rate=6e-4, gradient_accumulation_steps=4, gradient_checkpointing=True, dataloader_pin_memory=True, dataloader_num_workers=4, # torch_compile=True, # checkpointing: save_steps=370, optim="adafactor", # optim="adamw_torch", fp16=True, # push_to_hub=True, )

But I'm getting this error when running trainer.train();

[syncdoth]

Recent commit should solve this issue