fix: Subset Optuna

2025-12-10 11:09:44 +01:00 · 2025-12-10 11:09:44 +01:00 · 956ff79fa1
commit 956ff79fa1
parent 51e0ed7fc0
1 changed files with 6 additions and 1 deletions
--- a/main.py
+++ b/main.py
@ -12,12 +12,17 @@ def main():

    match args.mode:
        case 'train':
+            size = None
+            if args.method == 'optuna':
+                size = 2 ** 12
+            if args.debug:
+                size = 2 ** 10
            train(
                device=device,
                dataset=args.dataset,
                data_root=args.data_root,
                n_trials=3 if args.debug else None,
-                size=2 ** 10 if args.debug else None,
+                size=size,
                method=args.method,
                model_name=args.model,
                model_path=args.model_load_path,