Fix test failure on DeepSpeed (#29444)

* Fix test failure * use item
huggingface · May 14, 2024 · 2b8a613 · 2b8a613
1 parent 93e8af4
commit 2b8a613
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