add missing condition for first moment type as statement
Summary: Pull Request resolved:
@@ -165,6 +165,7 @@ class Adafactor(torch.optim.Optimizer):
state['RMS'] = 0
else:
+ 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)
