Update python/cugraph-equivariant/cugraph_equivariant/nn/tensor_produ…

…ct_conv.py Co-authored-by: Mario Geiger <[email protected]>
rapidsai · Jan 24, 2024 · 6dff47a · 6dff47a
1 parent 40b972c
commit 6dff47a
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py b/python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py
@@ -186,6 +186,7 @@ def forward(
             Reduction operator. Choose between "mean" and "sum".
 
         edge_envelope: torch.Tensor, optional
+            Optional, typically used as attenuation factors to fade out messages coming from nodes close to the cutoff distance used to create the graph. This is important to make the model smooth to the changes in node's coordinates.
             Edge envelope.
             Shape: (num_edges,)