Merge pull request VITA-Group#134 from choprahetarth/master

Fixed error VITA-Group#133
562628426 · May 4, 2022 · 595f258 · 595f258
2 parents 756db7f + e36dc2f
commit 595f258
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/predict.py b/predict.py
@@ -16,7 +16,7 @@
 class Predictor:
     def __init__(self, weights_path: str, model_name: str = ''):
         with open('config/config.yaml') as cfg:
-            config = yaml.load(cfg)
+            config = yaml.load(cfg, Loader=yaml.FullLoader)
         model = get_generator(model_name or config['model'])
         model.load_state_dict(torch.load(weights_path)['model'])
         self.model = model.cuda()