fix smooth lddt loss

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -2047,7 +2047,7 @@ def forward(
 
             smooth_lddt_loss = self.smooth_lddt_loss(
                 denoised_atom_pos,
-                normalized_atom_pos,
+                atom_pos_ground_truth,
                 atom_is_dna,
                 atom_is_rna,
                 coords_mask = atom_mask
@@ -2488,7 +2488,7 @@ def __init__(
         self,
         *,
         dim_single_inputs,
-        atompair_dist_bins: Float['d'],
+        atompair_dist_bins: Float[' d'],
         dim_single = 384,
         dim_pairwise = 128,
         num_plddt_bins = 50,

alphafold3_pytorch/attention.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import NamedTuple
+from typing import NamedTuple, Tuple
 
 import torch
 from torch import nn

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.0.33"
+version = "0.0.34"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }

tests/test_af3.py

@@ -428,6 +428,7 @@ def test_alphafold3():
         pde_labels = pde_labels,
         plddt_labels = plddt_labels,
         resolved_labels = resolved_labels,
+        diffusion_add_smooth_lddt_loss = True,
         return_loss_breakdown = True
     )
 
@@ -580,7 +581,6 @@ def test_alphafold3_with_packed_atom_repr():
 
     loss.backward()
 
-    print(residue_atom_lens)
     sampled_atom_pos = alphafold3(
         num_sample_steps = 16,
         atom_inputs = atom_inputs,