diff --git a/src/careamics/config/vae_algorithm_model.py b/src/careamics/config/vae_algorithm_model.py
index 1474a2951..0c0e01d76 100644
--- a/src/careamics/config/vae_algorithm_model.py
+++ b/src/careamics/config/vae_algorithm_model.py
@@ -41,17 +41,17 @@ class VAEAlgorithmConfig(BaseModel):
     algorithm: Literal["musplit", "denoisplit"]
 
     # NOTE: these are all configs (pydantic models)
-    loss: LVAELossConfig
     model: Union[LVAEModel, CustomModel] = Field(discriminator="architecture")
+    loss: Optional[LVAELossConfig] = None
     noise_model: Optional[MultiChannelNMConfig] = None
     noise_model_likelihood: Optional[NMLikelihoodConfig] = None
     gaussian_likelihood: Optional[GaussianLikelihoodConfig] = None
 
     # Optional fields
-    optimizer: OptimizerModel = OptimizerModel()
+    optimizer: Optional[OptimizerModel] = OptimizerModel()
     """Optimizer to use, defined in SupportedOptimizer."""
 
-    lr_scheduler: LrSchedulerModel = LrSchedulerModel()
+    lr_scheduler: Optional[LrSchedulerModel] = LrSchedulerModel()
 
     @model_validator(mode="after")
     def algorithm_cross_validation(self: Self) -> Self: