Fix momentum in non-nesterov case

ac74f345 · Simon Layton · 6d6f0bc2 · ac74f345
Commit ac74f345 authored Mar 08, 2019 by Simon Layton
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 0 deletions

csrc/multi_tensor_sgd_kernel.cu csrc/multi_tensor_sgd_kernel.cu +2 -0

No files found.
--- a/csrc/multi_tensor_sgd_kernel.cu
+++ b/csrc/multi_tensor_sgd_kernel.cu
@@ -104,6 +104,8 @@ struct SGDFunctor

            if (nesterov) {
              incoming_grads[ii] += momentum * incoming_moms[ii];
+            } else {
+              incoming_grads[ii] = incoming_moms[ii];
            }
          }