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.
A bug in scatter_add could cause a dimension mismatch in the
return z + self.bias
line ofSETLayer.forward
. The callscatter_add(k, self.inds_out)
produces a tensor of lengthmax(self.inds_out)+1
, which might not be equal toself.outdim
. This commit fixes this issue by using thedim_size
keyword of thescatter_add
function.