make sure all plms are excluded from state_dict

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -82,7 +82,8 @@
 
 from alphafold3_pytorch.plm import (
     PLMEmbedding,
-    PLMRegistry
+    PLMRegistry,
+    remove_plms
 )
 
 from alphafold3_pytorch.utils.model_utils import (
@@ -6242,6 +6243,14 @@ def __init__(
     def device(self):
         return self.zero.device
 
+    @remove_plms
+    def state_dict(self, *args, **kwargs):
+        return super().state_dict(*args, **kwargs)
+
+    @remove_plms
+    def load_state_dict(self, *args, **kwargs):
+        return super().load_state_dict(*args, **kwargs)
+
     @property
     def state_dict_with_init_args(self):
         return dict(

alphafold3_pytorch/plm.py

@@ -1,5 +1,5 @@
 import re
-from functools import partial
+from functools import partial, wraps
 
 import torch
 from beartype.typing import Literal
@@ -15,6 +15,22 @@
 def join(arr, delimiter = ''): # just redo an ugly part of python
     return delimiter.join(arr)
 
+def remove_plms(fn):
+    @wraps(fn)
+    def inner(self, *args, **kwargs):
+        has_plms = hasattr(self, 'plms')
+        if has_plms:
+            plms = self.plms
+            delattr(self, 'plms')
+
+        out = fn(self, *args, **kwargs)
+
+        if has_plms:
+            self.plms = plms
+
+        return out
+    return inner
+
 # constants
 
 aa_constants = get_residue_constants(res_chem_index=IS_PROTEIN)

tests/test_af3.py

@@ -1098,6 +1098,10 @@ def test_alphafold3_with_plm_embeddings():
         plm_embeddings="esm2_t33_650M_UR50D",
     )
 
+    state_dict = alphafold3.state_dict()
+
+    assert not any([key.startswith('plms') for key in state_dict.keys()])
+
     # mock inputs
 
     seq_len = 16