fix nan values

335525a7 · rusty1s · 68f4609c · 335525a7
Commit 335525a7 authored Jan 13, 2021 by rusty1s
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 0 deletions

csrc/scatter.cpp csrc/scatter.cpp +1 -0

No files found.
--- a/csrc/scatter.cpp
+++ b/csrc/scatter.cpp
@@ -97,6 +97,7 @@ public:
    auto dim = ctx->saved_data["dim"].toInt();
    auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
    auto grad_in = torch::gather(grad_out * out, dim, index, false).div_(src);
+    grad_in.masked_fill_(grad_in.isnan(), 0);
    return {grad_in, Variable(), Variable(), Variable(), Variable()};
  }
 };