diff --git a/alphafold3_pytorch/alphafold3.py b/alphafold3_pytorch/alphafold3.py
index 18f3c5db..ddd31070 100644
--- a/alphafold3_pytorch/alphafold3.py
+++ b/alphafold3_pytorch/alphafold3.py
@@ -5307,7 +5307,7 @@ def __init__(
             1.4      # 5 - water
         ])
 
-        self.register_buffer('atom_radii', atom_type_radii)
+        self.register_buffer('atom_radii', atom_type_radii, persistent = False)
 
     @property
     def is_mkdssp_available(self):