fix bugs in CPU adam (#633)

* add cpu adam counter for all cpu adam * fixed updating error in adam kernel
2025-09-13 05:01:44 +00:00 · 2022-04-02 17:04:05 +08:00
parent 1e2557e801
commit b31daed4cf
6 changed files with 25 additions and 13 deletions
--- a/colossalai/kernel/cuda_native/csrc/cpu_adam.cpp
+++ b/colossalai/kernel/cuda_native/csrc/cpu_adam.cpp
@@ -493,7 +493,7 @@ int adam_step(int optimizer_id,
                grads_ptr,
                exp_avg_ptr,
                exp_avg_sq_ptr,
-                params_c.size(0),
+                params_c.numel(),
                (params.options().dtype() == at::kHalf),
                (grads.options().dtype() == at::kHalf),
                loss_scale);