output type for distogram head

lucidrains · May 19, 2024 · b0198d4 · b0198d4
1 parent 216eea7
commit b0198d4
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/alphafold3_pytorch/alphafold3.py b/alphafold3_pytorch/alphafold3.py
@@ -1720,11 +1720,13 @@ def forward(
 # distogram head
 
 class DistogramHead(Module):
+
+    @typecheck
     def __init__(
         self,
         *,
         dim_pairwise = 128,
-        num_dist_bins = 38   # think it is 38?
+        num_dist_bins = 38,   # think it is 38?
     ):
         super().__init__()
 
@@ -1737,9 +1739,9 @@ def __init__(
     def forward(
         self,
         pairwise_repr: Float['b n n d']
-    ):
-        logits = self.to_distogram_logits(pairwise_repr)
+    ) -> Float['b l n n']:
 
+        logits = self.to_distogram_logits(pairwise_repr)
         return logits
 
 # confidence head