Add warning to test utils script

94d890d8 · Gustaf Ahdritz · 6e5c619c · 94d890d8
Commit 94d890d8 authored Nov 15, 2021 by Gustaf Ahdritz
Show whitespace changes
Inline Side-by-side

Showing with 7 additions and 1 deletion

tests/compare_utils.py tests/compare_utils.py +7 -1

No files found.
--- a/tests/compare_utils.py
+++ b/tests/compare_utils.py
@@ -100,5 +100,11 @@ def fetch_alphafold_module_weights(weight_path):
        module_name = spl[-1]
        prefix = "/".join(spl[:-1]) + "/"
        _remove_key_prefix(params, prefix)
+    try:
        params = alphafold.model.utils.flat_params_to_haiku(params)
+    except:
+        raise ImportError(
+            "Make sure to call import_alphafold before running this function"
+        )
    return params