diff --git a/python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py b/python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py index 998b7013446..ea9979427d0 100644 --- 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,)