add missing condition for first moment type as statement

Summary: Pull Request resolved: https://github.com/fairinternal/fairseq-py/pull/684

Differential Revision: D16006333

Pulled By: myleott

fbshipit-source-id: 95bd4215734281194008fa029e81407d63b335ac

fairseq/optim/adafactor.py

@@ -165,7 +165,8 @@ class Adafactor(torch.optim.Optimizer):
 
                     state['RMS'] = 0
                 else:
-                    state['exp_avg'] = state['exp_avg'].type_as(grad)
+                    if use_first_moment:
+                        state['exp_avg'] = state['exp_avg'].type_as(grad)
                     if factored:
                         state['exp_avg_sq_row'] = state['exp_avg_sq_row'].type_as(grad)
                         state['exp_avg_sq_col'] = state['exp_avg_sq_col'].type_as(grad)