fixing a bug in CPU Adam and Adagrad (#3109)

Co-authored-by: Bing Xie <bingxie@BINGHYPC014.redmond.corp.microsoft.com>
Co-authored-by: Shaden Smith <Shaden.Smith@microsoft.com>
Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com>
Co-authored-by: Jeff Rasley <jerasley@microsoft.com>

csrc/adagrad/cpu_adagrad.cpp

@@ -178,7 +178,7 @@ int ds_adagrad_step(int optimizer_id,
         std::static_pointer_cast<Adagrad_Optimizer>(s_optimizers[optimizer_id]);
     opt->IncrementStep(step);
     opt->update_state(lr, epsilon, weight_decay);
-    opt->Step_8(params_ptr, grads_ptr, exp_avg_sq_ptr, params_c.size(0));
+    opt->Step_8(params_ptr, grads_ptr, exp_avg_sq_ptr, params_c.numel());
 
 #if defined(__ENABLE_CUDA__)
     opt->SynchronizeStreams();
@@ -214,7 +214,7 @@ int ds_adagrad_step_plus_copy(int optimizer_id,
     opt->Step_8(params_ptr,
                 grads_ptr,
                 exp_avg_sq_ptr,
-                params_c.size(0),
+                params_c.numel(),
                 gpu_params_ptr,
                 (params.options().dtype() == at::kHalf));
 

csrc/adam/cpu_adam.cpp

@@ -235,7 +235,7 @@ int ds_adam_step(int optimizer_id,
                 grads_ptr,
                 exp_avg_ptr,
                 exp_avg_sq_ptr,
-                params_c.size(0),
+                params_c.numel(),
                 nullptr,
                 (params.options().dtype() == at::kHalf));
 
@@ -280,7 +280,7 @@ int ds_adam_step_plus_copy(int optimizer_id,
                 grads_ptr,
                 exp_avg_ptr,
                 exp_avg_sq_ptr,
-                params_c.size(0),
+                params_c.numel(),
                 gpu_params_ptr,
                 (params.options().dtype() == at::kHalf));