correct an update to predict.py

henryaddison · Feb 13, 2024 · c3227f3 · c3227f3
1 parent 4c0b665
commit c3227f3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/bin/predict.py b/bin/predict.py
@@ -120,7 +120,7 @@ def load_model(config, ckpt_filename):
     state = _init_state(config)
     state, loaded = restore_checkpoint(ckpt_filename, state, config.device)
     assert loaded, "Did not load state from checkpoint"
-    state["ema"].copy_to(state["score_model"].parameters())
+    state["ema"].copy_to(state["model"].parameters())
 
     # Sampling
     num_output_channels = len(get_variables(config.data.dataset_name)[1])