Bug Description

[A clear, concise description of the bug]

in optim/AdamOptimezer.cpp: the parameter count_ has an operation to add itself in foreach of parameters.
This means different parameters will use different counter in one iteration. it is in AdamOptimezer.cpp line 77, count_++ should be outside the for statement

void AdamOptimizer::step() {
// count_++ // it may be here outside for statement
for (size_t i = 0; i < parameters_.size(); i++) {
if (!parameters_[i].isGradAvailable()) {
continue;
}

const af::array& grad = parameters_[i].grad().array();
af::array& data = parameters_[i].array();

if (wd_ != 0) {
  // Weight decay term
  data = data - wd_ * data;
}

af::array& biasedFirst = biasedFirst_[i];
af::array& biasedSecond = biasedSecond_[i];

biasedFirst = beta1_ * biasedFirst + (1 - beta1_) * grad;
biasedSecond = beta2_ * biasedSecond + (1 - beta2_) * grad * grad;

af::eval(biasedFirst);
af::eval(biasedSecond);

//here, line 77: count_++
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);

}
}

Reproduction Steps

[Steps to reproduce. Please include, if needed for your issue:

• Any configurations or custom setup
• Commands to run
• A trace of the error]

Platform and Hardware

[Please list your operating system, [GPU] hardware, compiler, and other details if relevant]

Additional Context

[Add any additional information here]