X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=main.py;h=b9b52d6e18162292f09394de9d9150edf281a9f3;hb=6c8bed86221baae24a7c2aaaa41c009444efb5c9;hp=1b0d39a004436466724e13144599102e4e96b3a3;hpb=ef3bef5253ff719953dfffff28d4122c19acdd77;p=picoclvr.git

diff --git a/main.py b/main.py
index 1b0d39a..b9b52d6 100755
--- a/main.py
+++ b/main.py
@@ -12,7 +12,7 @@ from torch import nn
 from torch.nn import functional as F
 
 import ffutils
-import mygpt, tasks
+import mygpt, tasks, problems
 
 ######################################################################
 
@@ -335,19 +335,19 @@ picoclvr_pruner_eval = (
 
 if args.task == "sandbox":
     if args.sandbox_level == 0:
-        problem = tasks.ProblemLevel0(
+        problem = problems.ProblemLevel0(
             nb_sentences=args.sandbox_levels_nb_items,
             len_prompt=args.sandbox_levels_len_source,
             len_result=args.sandbox_levels_len_result,
         )
     elif args.sandbox_level == 1:
-        problem = tasks.ProblemLevel1(
+        problem = problems.ProblemLevel1(
             nb_operators=args.sandbox_levels_nb_items,
             len_source=args.sandbox_levels_len_source,
             len_result=args.sandbox_levels_len_result,
         )
     elif args.sandbox_level == 2:
-        problem = tasks.ProblemLevel2(
+        problem = problems.ProblemLevel2(
             len_source=args.sandbox_levels_len_source,
             len_result=args.sandbox_levels_len_result,
         )
@@ -355,8 +355,9 @@ if args.task == "sandbox":
         raise ValueError(f"Unknown sandbox level {args.sandbox_level}")
 
     task = tasks.SandBox(
-        problem,
-        # tasks.ProblemAddition(zero_padded=False, inverted_result=False),
+        # problem,
+        # problems.ProblemAddition(zero_padded=False, inverted_result=False),
+        problems.ProblemLenId(len_max=args.sandbox_levels_len_source),
         nb_train_samples=args.nb_train_samples,
         nb_test_samples=args.nb_test_samples,
         batch_size=args.batch_size,