Fix apex Adam to not break CPU mode
@@ -14,13 +14,17 @@ from . import FairseqOptimizer, register_optimizer
@register_optimizer('adam')
class FairseqAdam(FairseqOptimizer):
+
def __init__(self, args, params):
super().__init__(args, params)
+ if torch.cuda.is_available():
try:
from apex.optimizers import FusedAdam
self._optimizer = FusedAdam(params, **self.optimizer_config)
except ImportError:
self._optimizer = Adam(params, **self.optimizer_config)
+ else:
+ self._optimizer = Adam(params, **self.optimizer_config)
@staticmethod
def add_args(parser):
