Implemented AtomAttentionEncoder and AtomAttentionDecoder, as well as…

… token-to-atom mapping functions
Ligo-Biosciences · May 29, 2024 · d96d634 · d96d634
1 parent 2474b04
commit d96d634
Show file tree

Hide file tree

Showing 4 changed files with 584 additions and 67 deletions.
diff --git a/src/diffusion/attention.py b/src/diffusion/attention.py
@@ -28,8 +28,8 @@ def __init__(
             num_blocks:
                 Number of blocks.
             num_heads:
-                Number of parallel attention heads. Note that embed_dim will be split across num_heads
-                (i.e. each head will have dimension embed_dim // num_heads).
+                Number of parallel attention heads. Note that c_atom will be split across num_heads
+                (i.e. each head will have dimension c_atom // num_heads).
             dropout:
                 Dropout probability on attn_output_weights. Default: 0.0 (no dropout).
         """