-def generate(nb_values=3, max_input=9, prog_len=6, nb_runs=5):
- prog_len = 1 + torch.randint(prog_len - 1, (1,)).item()
+def generate(nb_starting_values=3, max_input=9, prog_len=6, nb_runs=5):
+ prog_len = (1 + torch.randint(2 * prog_len, (1,))).clamp(max=prog_len).item()