Use shimmer migrations in migrate script

ruflab · Mar 25, 2024 · 97658e7 · 97658e7
1 parent 5006441
commit 97658e7
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/playground/migrate_model.py b/playground/migrate_model.py
@@ -1,3 +1,5 @@
+from shimmer import migrate_model as migrate_shimmer_model
+
 from simple_shapes_dataset import DEBUG_MODE, PROJECT_DIR
 from simple_shapes_dataset.ckpt_migrations import migrate_model
 from simple_shapes_dataset.config import load_config
@@ -10,6 +12,7 @@ def main():
     )
 
     if config.global_workspace.checkpoint is not None:
+        migrate_shimmer_model(config.global_workspace.checkpoint)
         migrate_model(
             config.global_workspace.checkpoint, PROJECT_DIR / "migrations" / "gw"
         )