FEA:add SparseAdam trainer

RUCAIBox · Dec 10, 2020 · 05e9f56 · 05e9f56
1 parent 1789405
commit 05e9f56
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/recbole/trainer/trainer.py b/recbole/trainer/trainer.py
@@ -116,6 +116,8 @@ def _build_optimizer(self):
             optimizer = optim.Adagrad(self.model.parameters(), lr=self.learning_rate)
         elif self.learner.lower() == 'rmsprop':
             optimizer = optim.RMSprop(self.model.parameters(), lr=self.learning_rate)
+        elif self.learner.lower() == 'sparse_adam':
+            optimizer =optim.SparseAdam(self.model.parameters(), lr=self.learning_rate)
         else:
             self.logger.warning('Received unrecognized optimizer, set default Adam optimizer')
             optimizer = optim.Adam(self.model.parameters(), lr=self.learning_rate)