remove atomic

Author: eric-haibin-lin

src/operator/contrib/multi_lamb-inl.h

@@ -302,10 +302,10 @@ inline void multiLAMB(const nnvm::NodeAttrs& attrs,
     Tensor<xpu, 1, int> block_to_chunk(reinterpret_cast<int*>(&workspace[pos_wspace]),
       Shape1(kernel_params.nchunks), s);
 
-    MultiSumSqRun<xpu>(weights, kernel_params.ntensors, r1.dptr_, s);
+    MultiSumSqRun<xpu>(weights, kernel_params.ntensors, r1.dptr_, ctx);
     call_kernel1<MPDType, DType>(s, kernel_params, param, block_to_tensor.dptr_,
                                   block_to_chunk.dptr_);
-    MultiSumSqRun<xpu>(temp_g, kernel_params.ntensors, r2.dptr_, s);
+    MultiSumSqRun<xpu>(temp_g, kernel_params.ntensors, r2.dptr_, ctx);
     call_kernel2<MPDType, DType>(s, kernel_params, param, r1.dptr_, r2.dptr_,
                                  block_to_tensor.dptr_, block_to_chunk.dptr_,
                                  req[0]);