- train_set = CompressedVignetteSet(problem_number, args.nb_train_batches, args.batch_size)
- test_set = CompressedVignetteSet(problem_number, args.nb_test_batches, args.batch_size)
+ train_set = CompressedVignetteSet(problem_number, args.nb_train_batches, args.batch_size,
+ cuda=torch.cuda.is_available())
+ test_set = CompressedVignetteSet(problem_number, args.nb_test_batches, args.batch_size,
+ cuda=torch.cuda.is_available())