bugfixes and adds a section to convert optim files

openfold/utils/import_weights.py

@@ -670,6 +670,7 @@ def import_jax_weights_(model, npz_path, version="model_1"):
 
 def convert_deprecated_v1_keys(state_dict):
     """Update older OpenFold model weight names to match the current model code."""
+    logging.warning('converting keys...')
 
     replacements = {
         'template_angle_embedder': 'template_single_embedder',

train_openfold.py

@@ -43,6 +43,7 @@ from scripts.zero_to_fp32 import (
     get_fp32_state_dict_from_zero_checkpoint,
     get_global_step_from_zero_checkpoint
 )
+from scripts.zero_to_fp32 import get_optim_files, parse_optim_states, get_model_state_file
 
 from openfold.utils.logger import PerformanceLoggingCallback