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

huggingface · Sep 2, 2022 · 9b3eb81 · 9b3eb81
1 parent 142e12a
commit 9b3eb81
Showing 1 changed file with 2 additions and 0 deletions.
diff --git 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