bugfix: Catch learn-mask error when no mask type selected

deepfakes · Oct 12, 2022 · 403f981 · 403f981
1 parent f1e3339
commit 403f981
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/plugins/train/model/_base/model.py b/plugins/train/model/_base/model.py
@@ -129,6 +129,10 @@ def __init__(self,
                                 "Mask to use. Please select a mask or disable Penalized Mask "
                                 "Loss.")
 
+        if self.config["learn_mask"] and self.config["mask_type"] is None:
+            raise FaceswapError("'Learn Mask' has been selected but you have not chosen a Mask to "
+                                "use. Please select a mask or disable 'Learn Mask'.")
+
         self._mixed_precision = self.config["mixed_precision"] and get_backend() != "amd"
         # self._io = IO(self, model_dir, self._is_predict, self.config["save_optimizer"])
         # TODO - Re-enable saving of optimizer once this bug is fixed: