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_samples)
test_input, test_target = generate_set(problem_number, args.nb_test_samples)