Revert _MultiHeadAttention to MultiHeadAttention

torchmd · Oct 30, 2023 · 237f4b4 · 237f4b4
1 parent 6c98d69
commit 237f4b4
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/torchmdnet/models/torchmd_t.py b/torchmdnet/models/torchmd_t.py
@@ -115,7 +115,7 @@ def __init__(
 
         self.attention_layers = nn.ModuleList()
         for _ in range(num_layers):
-            layer = _MultiHeadAttention(
+            layer = MultiHeadAttention(
                 hidden_channels,
                 num_rbf,
                 distance_influence,
@@ -182,7 +182,7 @@ def __repr__(self):
         )
 
 
-class _MultiHeadAttention(MessagePassing):
+class MultiHeadAttention(MessagePassing):
     def __init__(
         self,
         hidden_channels,
@@ -195,7 +195,7 @@ def __init__(
         cutoff_upper,
         dtype=torch.float,
     ):
-        super(_MultiHeadAttention, self).__init__(aggr="add", node_dim=0)
+        super(MultiHeadAttention, self).__init__(aggr="add", node_dim=0)
         assert hidden_channels % num_heads == 0, (
             f"The number of hidden channels ({hidden_channels}) "
             f"must be evenly divisible by the number of "