[BUG] Setting average_tokens_across_devices to True caused an error because it attempted to gather CPU tensors using NCCL.

[Jintao-Huang]

transformers/src/transformers/trainer.py

Lines 5288 to 5297 in 1dfad4b

	if num_items_in_batch is not None:
	if self.args.average_tokens_across_devices:
	num_items_in_batch = self.accelerator.gather(num_items_in_batch).sum()
	if torch.is_tensor(num_items_in_batch):
	num_items_in_batch = num_items_in_batch.to(device)
	if self.args.n_gpu > 1 and num_items_in_batch.dim() == 0:
	# In the DataParallel case, convert the scalar tensor into a 1-dim tensor
	num_items_in_batch = num_items_in_batch.unsqueeze(0)

[Jintao-Huang]

[Rocketknight1]

cc @SunMarc

[SunMarc]

Hey @Jintao-Huang, can you share a reproducer ? The num_items_in_batch should be on cuda normally