X-Git-Url: https://www.fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=main.py;h=3bf7587a5f1e45bee2530ca56ce016738b95d117;hb=046f35f38d629c9854104e855a53f0142449138f;hp=a31284e68ecf616231d89c27d324403c1ab45288;hpb=9425d322d40205b07000fd1bd23ef0696085c30b;p=mygpt.git

diff --git a/main.py b/main.py
index a31284e..3bf7587 100755
--- a/main.py
+++ b/main.py
@@ -111,12 +111,20 @@ import picoclvr
 
 class TaskPicoCLVR(Task):
 
-    def __init__(self, batch_size, height = 6, width = 8, device = torch.device('cpu')):
+    def __init__(self, batch_size,
+                 height = 6, width = 8, many_colors = False,
+                 device = torch.device('cpu')):
+
         self.batch_size = batch_size
         self.device = device
         nb = args.data_size if args.data_size > 0 else 250000
 
-        descr = picoclvr.generate(nb, height = height, width = width)
+        descr = picoclvr.generate(
+            nb,
+            height = height, width = width,
+            many_colors = many_colors
+        )
+
         descr = [ s.strip().split(' ') for s in descr ]
         l = max([ len(s) for s in descr ])
         descr = [ s + [ '<unk>' ] * (l - len(s)) for s in descr ]