X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=main.py;h=958dfc70a791d2b47708b10e7f44e53e7456fc2e;hb=545b06257dd69285c6321d7bd713d819da74953b;hp=55f2c2f4c6f9347cc6b8d04d5311bab811e4057a;hpb=a3f5378c684fb58b0dc839638e768c3a4b8e8a83;p=picoclvr.git

diff --git a/main.py b/main.py
index 55f2c2f..958dfc7 100755
--- a/main.py
+++ b/main.py
@@ -89,7 +89,9 @@ parser.add_argument("--checkpoint_name", type=str, default="checkpoint.pth")
 ##############################
 # filetask
 
-parser.add_argument("--filetask_file", type=str, default=None)
+parser.add_argument("--filetask_train_file", type=str, default=None)
+
+parser.add_argument("--filetask_test_file", type=str, default=None)
 
 ##############################
 # rpl options
@@ -403,13 +405,15 @@ picoclvr_pruner_eval = (
 
 if args.task == "file":
     assert (
-        args.filetask_file is not None
-    ), "You have to specify the task file with --filetask_file <filename>"
+        args.filetask_train_file is not None and args.filetask_test_file is not None
+    ), "You have to specify the task train and test files"
     task = tasks.TaskFromFile(
-        args.filetask_file,
+        args.filetask_train_file,
+        args.filetask_test_file,
         nb_train_samples=args.nb_train_samples,
         nb_test_samples=args.nb_test_samples,
         batch_size=args.batch_size,
+        shuffle=True,
         device=device,
     )
     args.max_percents_of_test_in_train = 0