- train_input, train_target = generate_set(problem_number, args.nb_train_samples)
- test_input, test_target = generate_set(problem_number, args.nb_test_samples)
+ train_input, train_target = generate_set(problem_number,
+ args.nb_train_batches * args.batch_size)
+ test_input, test_target = generate_set(problem_number,
+ args.nb_test_batches * args.batch_size)