forgot loading checkpoint (issue #11)

ivanvovk · Nov 12, 2020 · d230621 · d230621
1 parent f59d4bd
commit d230621
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/inference.py b/inference.py
@@ -44,6 +44,7 @@
 
     # Initialize the model
     model = WaveGrad(config)
+    model.load_state_dict(torch.load(args.checkpoint_path)['model'], strict=False)
 
     # Set noise schedule
     noise_schedule = torch.load(args.noise_schedule_path)