Update main.py

add rollback to BSC option
tvlearn · Jan 25, 2024 · 8f3ff73 · 8f3ff73
1 parent e70a17a
commit 8f3ff73
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/examples/gaussian-denoising/main.py b/examples/gaussian-denoising/main.py
@@ -141,6 +141,7 @@ def gaussian_denoising_example():  # noqa: C901
         reco_epochs=reco_epochs,
         log_blacklist=[],
         log_only_latest_theta=True,
+        rrollback_if_F_decreases=['W', 'sigma2'] if args.model='bsc' else []
     )
     exp = Training(conf=exp_config, estep_conf=estep_conf, model=model, train_data_file=data_file)
     logger, trainer = exp.logger, exp.trainer