Fix accuracy error of NPUFusedAdam (microsoft#5777)

Co-authored-by: gp513 <[email protected]> Co-authored-by: Logan Adams <[email protected]>
samadejacobs · Jul 23, 2024 · 5a100f6 · 5a100f6
1 parent acdf136
commit 5a100f6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/op_builder/npu/fused_adam.py b/op_builder/npu/fused_adam.py
@@ -16,8 +16,8 @@ class NPUFusedAdam:
     @staticmethod
     def multi_tensor_adam(chunk_size, noop_flag_buffer, tensor_lists, lr, beta1, beta2, epsilon, step, adam_w_mode,
                           bias_correction, weight_decay, *args):
-        bias_correction1 = beta1**step
-        bias_correction2 = beta2**step
+        bias_correction1 = beta1**(step - 1)
+        bias_correction2 = beta2**(step - 1)
 
         # iteration group['params']
         for i in range(len(tensor_lists[0])):