Fix distillation bug

README.md

@@ -208,6 +208,10 @@ and supports the full range of training options that entails, including
 multi-node distributed training. For more information, consult PyTorch 
 Lightning documentation and the `--help` flag of the training script.
 
+Note that the data directory can also contain PDB files previously output by
+the model. These are treated as members of the self-distillation set and are
+subjected to distillation-set-only preprocessing steps.
+
 ## Testing
 
 To run unit tests, use

openfold/data/data_transforms.py

@@ -176,7 +176,7 @@ def randomly_replace_msa_with_unknown(protein, replace_proportion):
 
 @curry1
 def sample_msa(protein, max_seq, keep_extra, seed=None):
-    """Sample MSA randomly, remaining sequences are stored are stored as `extra_*`."""
+    """Sample MSA randomly, remaining sequences are stored are stored as `extra_*`.""" 
     num_seq = protein["msa"].shape[0]
     g = torch.Generator(device=protein["msa"].device)
     if seed is not None:
@@ -202,7 +202,7 @@ def sample_msa(protein, max_seq, keep_extra, seed=None):
 @curry1
 def sample_msa_distillation(protein, max_seq):
     if(protein["is_distillation"] == 1):
-        protein = sample_msa(protein, max_seq, keep_extra=False)
+        protein = sample_msa(max_seq, keep_extra=False)(protein)
     return protein