From fbf198b1595166a96c16863868b497a6485e1863 Mon Sep 17 00:00:00 2001
From: Samuel Larkin <Samuel.Larkin@cnrc-nrc.gc.ca>
Date: Fri, 1 Nov 2024 12:05:07 -0400
Subject: [PATCH] fix: call check_and_upgrade_checkpoint()

---
 hfgl/utils.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/hfgl/utils.py b/hfgl/utils.py
index 2e88274..b400478 100644
--- a/hfgl/utils.py
+++ b/hfgl/utils.py
@@ -42,9 +42,12 @@ def load_hifigan_from_checkpoint(ckpt: dict, device) -> Tuple[HiFiGAN, HiFiGANCo
         model = HiFiGAN(config).to(device)
     else:
         model = HiFiGANGenerator(config).to(device)
+
+    ckpt = model.check_and_upgrade_checkpoint(ckpt)
     model.load_state_dict(ckpt["state_dict"])
     model.generator.eval()
     model.generator.remove_weight_norm()
+
     return model, config