fix of adam optimizer as pointed in #143

Summary: fix for #143

Reviewed By: syhw, jacobkahn

Differential Revision: D22757150

fbshipit-source-id: 1c9267a7a594f7d3b7041389b7e40715cb14e6b4

flashlight/optim/AdamOptimizer.cpp

@@ -52,6 +52,11 @@ AdamOptimizer::AdamOptimizer(
 }
 
 void AdamOptimizer::step() {
+  count_++;
+  float correctedBias1 = 1 - std::pow(beta1_, count_);
+  float correctedBias2 = 1 - std::pow(beta2_, count_);
+  float correctedLr = lr_ * std::sqrt(correctedBias2) / correctedBias1;
+
   for (size_t i = 0; i < parameters_.size(); i++) {
     if (!parameters_[i].isGradAvailable()) {
       continue;
@@ -74,12 +79,6 @@ void AdamOptimizer::step() {
     af::eval(biasedFirst);
     af::eval(biasedSecond);
 
-    count_++;
-
-    float correctedBias1 = 1 - std::pow(beta1_, count_);
-    float correctedBias2 = 1 - std::pow(beta2_, count_);
-    float correctedLr = lr_ * std::sqrt(correctedBias2) / correctedBias1;
-
     data = data - (correctedLr * biasedFirst) / (af::sqrt(biasedSecond) + eps_);
 
     af::eval(data);