use in-place exp() + remove eps

torch_scatter/composite/softmax.py

@@ -4,8 +4,8 @@ from torch_scatter import scatter_sum, scatter_max
 from torch_scatter.utils import broadcast
 
 
-def scatter_softmax(src: torch.Tensor, index: torch.Tensor, dim: int = -1,
-                    eps: float = 1e-12) -> torch.Tensor:
+def scatter_softmax(src: torch.Tensor, index: torch.Tensor,
+                    dim: int = -1) -> torch.Tensor:
     if not torch.is_floating_point(src):
         raise ValueError('`scatter_softmax` can only be computed over tensors '
                          'with floating point data types.')
@@ -16,10 +16,10 @@ def scatter_softmax(src: torch.Tensor, index: torch.Tensor, dim: int = -1,
     max_per_src_element = max_value_per_index.gather(dim, index)
 
     recentered_scores = src - max_per_src_element
-    recentered_scores_exp = recentered_scores.exp()
+    recentered_scores_exp = recentered_scores.exp_()
 
     sum_per_index = scatter_sum(recentered_scores_exp, index, dim)
-    normalizing_constants = sum_per_index.add_(eps).gather(dim, index)
+    normalizing_constants = sum_per_index.gather(dim, index)
 
     return recentered_scores_exp.div(normalizing_constants)