if learning rate is a tensor, get item (float) (#18861)

9b3eb810 · kmckiern · GitHub · 142e12af · 9b3eb810
Unverified Commit 9b3eb810 authored Sep 02, 2022 by kmckiern Committed by GitHub Sep 02, 2022
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 0 deletions

src/transformers/trainer_pt_utils.py src/transformers/trainer_pt_utils.py +2 -0

No files found.
--- a/src/transformers/trainer_pt_utils.py
+++ b/src/transformers/trainer_pt_utils.py
@@ -837,6 +837,8 @@ def _get_learning_rate(self):
            if version.parse(version.parse(torch.__version__).base_version) >= version.parse("1.4")
            else self.lr_scheduler.get_lr()[0]
        )
+        if torch.is_tensor(last_lr):
+            last_lr = last_lr.item()
    return last_lr