From 1d09a20fc1b48a71edc3e0ae309f4ff704554b7c Mon Sep 17 00:00:00 2001 From: Aman Gupta Karmani Date: Mon, 2 Sep 2024 08:43:18 -0700 Subject: [PATCH] use duck-typing to ensure underlying optimizer supports schedulefree hooks (#3055) * use duck-typing to ensure underlying optimizer supports schedulefree hooks * fixup --- src/accelerate/optimizer.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/accelerate/optimizer.py b/src/accelerate/optimizer.py index 3230aa6606e..1c0a777dcba 100644 --- a/src/accelerate/optimizer.py +++ b/src/accelerate/optimizer.py @@ -125,13 +125,15 @@ def train(self): """ Sets the optimizer to "train" mode. Useful for optimizers like `schedule_free` """ - return self.optimizer.train() + if hasattr(self.optimizer, "train") and callable(self.optimizer.train): + self.optimizer.train() def eval(self): """ Sets the optimizer to "eval" mode. Useful for optimizers like `schedule_free` """ - return self.optimizer.eval() + if hasattr(self.optimizer, "eval") and callable(self.optimizer.eval): + self.optimizer.eval() def step(self, closure=None): if is_lomo_available():