Editing issue with pickle def with lambda function (huggingface#23869)

* Editing issue with pickle def with lambda function * fix type * Made helper function private * delete tab --------- Co-authored-by: georgebredis <[email protected]>
gojiteji · Jun 5, 2023 · a9fc1ea · a9fc1ea
1 parent 576755e
commit a9fc1ea
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/transformers/optimization.py b/src/transformers/optimization.py
@@ -32,6 +32,10 @@
 logger = logging.get_logger(__name__)
 
 
+def _get_constant_lambda(_=None):
+    return 1
+
+
 def get_constant_schedule(optimizer: Optimizer, last_epoch: int = -1):
     """
     Create a schedule with a constant learning rate, using the learning rate set in optimizer.
@@ -46,7 +50,7 @@ def get_constant_schedule(optimizer: Optimizer, last_epoch: int = -1):
         `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
     """
 
-    return LambdaLR(optimizer, lambda _: 1, last_epoch=last_epoch)
+    return LambdaLR(optimizer, _get_constant_lambda, last_epoch=last_epoch)
 
 
 def get_reduce_on_plateau_schedule(optimizer: Optimizer):