Fix test failure on DeepSpeed (#29444)

* Fix test failure * use item
huggingface · Mar 6, 2024 · 9322576 · 9322576
1 parent 0a5b051
commit 9322576
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
@@ -2011,7 +2011,10 @@ def _inner_training_loop(
  is_accelerate_available()
  and self.accelerator.distributed_type == DistributedType.DEEPSPEED
  ):
- grad_norm = model.get_global_grad_norm().item()
+ grad_norm = model.get_global_grad_norm()
+ # In some cases the grad norm may not return a float
+ if hasattr(grad_norm, "item"):
+ grad_norm = grad_norm.item()
  else:
  grad_norm = _grad_norm.item() if _grad_norm is not None else None