Fix test failure on DeepSpeed (#29444)

* Fix test failure * use item

Fix test failure on DeepSpeed (#29444)
* Fix test failure * use item
9322576e · Zach Mueller · GitHub · 0a5b0516 · 9322576e
Unverified Commit 9322576e authored Mar 06, 2024 by Zach Mueller Committed by GitHub Mar 06, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

src/transformers/trainer.py src/transformers/trainer.py +4 -1

No files found.
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -2011,7 +2011,10 @@ class Trainer:
                            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