Update pseudo_moe_st_module.py

pseudo_moe_st_module.py

@@ -84,9 +84,10 @@ class PseudoMoETransformer(nn.Module):
                 )
             domain = self.default_domain
         else:
-            if domain not in self.config.domain_names:
+            # "general" (the base model without LoRA) is always valid but is not listed in config.domain_names.
+            if domain != "general" and domain not in self.config.domain_names:
                 raise ValueError(
-                    f"Invalid domain: {domain}. Must be one of {self.config.domain_names}."
+                    f"Invalid domain: {domain}. Must be one of {['general'] + list(self.config.domain_names)}."
                 )
         self.set_domain(domain)
         
@@ -130,4 +131,4 @@ class PseudoMoETransformer(nn.Module):
 
     @property
     def domain(self) -> str:
-        return self.auto_model.domain
+        return self.auto_model.domain