Merge pull request dmlc#115 from sbodenstein/master

Fix bug in AddTakeGrad
hughperkins · Apr 25, 2016 · f08c7b6 · f08c7b6
2 parents be90db1 + 2cce4d5
commit f08c7b6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mshadow/cuda/tensor_gpu-inl.cuh b/mshadow/cuda/tensor_gpu-inl.cuh
@@ -473,7 +473,7 @@ inline void AddTakeGrad(Tensor<gpu, 2, DType> dst,
                         const Tensor<gpu, 2, DType> &src) {
   const int kUnitBits = kMemUnitBits + 1;
   dim3 dimBlock(1 << kUnitBits);
-  dim3 dimGrid(dst.size(1) >> kUnitBits);
+  dim3 dimGrid((dst.size(1) + (1 << kUnitBits) - 1) >> kUnitBits);
 
   CHECK_EQ(dst.size(1), src.size(1)) << "AddtTakeGrad: shape mismatch";
   CHECK_EQ(index.size(0), src.size(0)) << "AddtTakeGrad: shape mismatch";