fix

main.py

@@ -289,7 +289,7 @@ args.shared_group_mlp = not args.normal_group_mlp
 # And w=0 and cpu=2 for symbolic inputs
 if args.num_workers is None:
     if args.use_visual_inputs:
-        args.num_workers = 2
+        args.num_workers = 16
     else:
         args.num_workers = 0
 
@@ -467,7 +467,7 @@ def main():
     for epoch in range(args.epochs):
         train(trainloader, optimizer, model, epoch)
         validate(validloader, model, epoch)
-        test(trainloader, model, epoch)
+        test(testloader, model, epoch)
 
 if __name__ == '__main__':
     main()