Update tensornet.py

update ref
torchmd · Dec 17, 2023 · a06d072 · a06d072
1 parent af1aa23
commit a06d072
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/torchmdnet/models/tensornet.py b/torchmdnet/models/tensornet.py
@@ -56,8 +56,9 @@ def tensor_norm(tensor):
     return (tensor**2).sum((-2, -1))
 
 class TensorNet(nn.Module):
-    r"""TensorNet's architecture.
-    From TensorNet: Cartesian Tensor Representations for Efficient Learning of Molecular Potentials; G. Simeon and G. de Fabritiis.
+    r"""TensorNet's architecture. From
+    TensorNet: Cartesian Tensor Representations for Efficient Learning of Molecular Potentials; G. Simeon and G. de Fabritiis.
+    NeurIPS 2023.
 
     Args:
         hidden_channels (int, optional): Hidden embedding size.