[CLEANUP]

alphafold3/model.py

@@ -596,11 +596,11 @@ def forward(
         """
         # Recycle bins
         # recyle_bins = []
-        
+
         # TODO: Input
         # TODO: Template
         # TODO: MSA
-        
+
         b, n, n_two, dim = pair_representation.shape
         b_two, n_two, dim_two = single_representation.shape
 
@@ -612,21 +612,18 @@ def forward(
         # Apply the 48 blocks of PairFormer
         x = self.pairformer(x)
         print(x.shape)
-        
+
         # Add the embeddings to the recycle bins
         # recyle_bins.append(x)
-
 
         # Diffusion
         x = self.diffuser(x, ground_truth)
-        
+
         # If return_confidence is True, return the confidence
         if return_confidence is True:
             x = self.confidence_projection(x)
             return x
-        
+
         # If return_loss is True, return the loss
         if return_embeddings is True:
             return x
-
-

model_example.py

@@ -1,8 +1,10 @@
-import torch 
+import torch
 from alphafold3 import AlphaFold3
 
 # Create random tensors
-x = torch.randn(1, 5, 5, 64)  # Shape: (batch_size, seq_len, seq_len, dim)
+x = torch.randn(
+    1, 5, 5, 64
+)  # Shape: (batch_size, seq_len, seq_len, dim)
 y = torch.randn(1, 5, 64)  # Shape: (batch_size, seq_len, dim)
 
 # Initialize AlphaFold3 model