parser.add_argument("--seed", type=int, default=0)
+parser.add_argument("--max_percents_of_test_in_train", type=int, default=1)
+
+########################################
+
parser.add_argument("--nb_epochs", type=int, default=None)
parser.add_argument("--batch_size", type=int, default=None)
parser.add_argument("--learning_rate_schedule", type=str, default="10: 2e-5,30: 4e-6")
+########################################
+
parser.add_argument("--model", type=str, default="37M")
parser.add_argument("--dim_model", type=int, default=None)
parser.add_argument("--dropout", type=float, default=0.1)
+########################################
+
parser.add_argument("--deterministic_synthesis", action="store_true", default=False)
parser.add_argument("--no_checkpoint", action="store_true", default=False)
task = tasks.SandBox(
# problem,
# problems.ProblemAddition(zero_padded=False, inverted_result=False),
- problems.ProblemLenId(len_max=args.sandbox_levels_len_source),
+ # problems.ProblemLenId(len_max=args.sandbox_levels_len_source),
+ problems.ProblemTwoTargets(len_total=16, len_targets=4),
nb_train_samples=args.nb_train_samples,
nb_test_samples=args.nb_test_samples,
batch_size=args.batch_size,
)
assert (
- nb_in_train <= nb_test // 100
-), "More than 1% of test samples are in the train set"
+ nb_in_train <= args.max_percents_of_test_in_train * nb_test / 100
+), f"More than {args.max_percents_of_test_in_train}% of test samples are in the train set"
##############################