distogram logits can now be atom resolution

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -5353,7 +5353,7 @@ def forward(
     ) -> (
         Float['b m 3'] |
         Tuple[Float['b m 3'], ConfidenceHeadLogits] |
-        Tuple[Float['b m 3'], ConfidenceHeadLogits, Float['b l n n']] |
+        Tuple[Float['b m 3'], ConfidenceHeadLogits, Float['b l n n'] | Float['b l m m']] |
         Float[''] |
         Tuple[Float[''], LossBreakdown]
     ):