fix nan values

rusty1s · rusty1s · commit 335525a745a9 · 2021-01-13T09:43:04.000+01:00
diff --git a/csrc/scatter.cpp b/csrc/scatter.cpp
@@ -97,6 +97,7 @@ class ScatterMul : public torch::autograd::Function<ScatterMul> {
     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()};
   }
 };

Original file line number	Diff line number	Diff line change
`@@ -97,6 +97,7 @@ class ScatterMul : public torch::autograd::Function<ScatterMul> {`
`97`	`97`	`auto dim = ctx->saved_data["dim"].toInt();`
`98`	`98`	`auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());`
`99`	`99`	`auto grad_in = torch::gather(grad_out * out, dim, index, false).div_(src);`
	`100`	`+ grad_in.masked_fill_(grad_in.isnan(), 0);`
`100`	`101`	`return {grad_in, Variable(), Variable(), Variable(), Variable()};`
`101`	`102`	`}`
`102`	`103`	`};`