From 1ac7dd66984ceeed9341381c6eda29ba7407e8bd Mon Sep 17 00:00:00 2001 From: rig1 Date: Mon, 27 Jan 2020 15:06:18 -0600 Subject: [PATCH] Fix scatter_add bug in SETLayer.forward A bug in scatter_add could cause a dimension mismatch in the `return z + self.bias` line of `SETLayer.forward`. The call `scatter_add(k, self.inds_out)` produces a tensor of length `max(self.inds_out)+1`, which might not be equal to `self.outdim`. This commit fixes this issue by using the `dim_size` keyword of the `scatter_add` function. --- synapses/SET_layer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/synapses/SET_layer.py b/synapses/SET_layer.py index 259996c..915c952 100644 --- a/synapses/SET_layer.py +++ b/synapses/SET_layer.py @@ -258,5 +258,5 @@ def reset_parameters(self): def forward(self, x): k = x[:, self.inds] k = k * self.weight - z = scatter_add(k, self.inds_out) + z = scatter_add(k, self.inds_out, dim_size=self.outdim) return z + self.bias \ No newline at end of file