         let learningRate = self.learningRate * 1 / (1 + decay * step)
         for kp in model.recursivelyAllWritableKeyPaths(to: Tensor<Scalar>.self) {
             alpha[keyPath: kp] =
-                rho * alpha[keyPath: kp] + (1 - rho) * pow(gradient[keyPath: kp], 2)
+                rho * alpha[keyPath: kp] + (1 - rho) * pow(vector[keyPath: kp], 2)
             model[keyPath: kp] -=
-                learningRate * gradient[keyPath: kp] / (sqrt(alpha[keyPath: kp]) + epsilon)
+                learningRate * vector[keyPath: kp] / (sqrt(alpha[keyPath: kp]) + epsilon)
         }
     }
 }


    public func update(_ model: inout Model.AllDifferentiableVariables,
                       along vector: Model.CotangentVector) {
        model = model.moved(along: learningRate * (.zero - model.tangentVector(from: gradient)))
    }
}

        let stepSize = learningRate * (sqrt(1 - pow(beta2, step)) / (1 - pow(beta1, step)))
        for kp in model.recursivelyAllWritableKeyPaths(to: Tensor<Scalar>.self) {
            firstMoments[keyPath: kp] =
                firstMoments[keyPath: kp] * beta1 + (1 - beta1) * vector[keyPath: kp]
            secondMoments[keyPath: kp] =
                secondMoments[keyPath: kp] * beta2 + (1 - beta2) *
                     vector[keyPath: kp] * vector[keyPath: kp]
            model[keyPath: kp] -=
                stepSize * firstMoments[keyPath: kp] / (sqrt(secondMoments[keyPath: kp]) + epsilon)
        }
