fix training

IAHispano · Sep 14, 2024 · 4a2ffbf · 4a2ffbf
1 parent 95c6677
commit 4a2ffbf
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/rvc/train/train.py b/rvc/train/train.py
@@ -134,7 +134,10 @@ def verify_checkpoint_shapes(checkpoint_path, model):
     checkpoint = torch.load(checkpoint_path, map_location="cpu")
     checkpoint_state_dict = checkpoint["model"]
     try:
-        model_state_dict = model.module.load_state_dict(checkpoint_state_dict)
+        if hasattr(model, 'module'):
+            model_state_dict = model.module.load_state_dict(checkpoint_state_dict)
+        else:
+            model_state_dict = model.load_state_dict(checkpoint_state_dict)
     except RuntimeError:
         print("The sample rate of the pretrain doesn't match the selected one")
         sys.exit(1)