Update app.py

app.py

@@ -206,7 +206,7 @@ def Generate_Rock_Song(input_midi,
 
     def generate_continuation(num_prime_tokens, num_gen_tokens):
 
-        x = torch.tensor(prime_toks[:num_prime_tokens], dtype=torch.long, device=DEVICE)
+        x = torch.tensor(TMIDIX.flatten(prime_toks)[:num_prime_tokens], dtype=torch.long, device=DEVICE)
         
         with ctx:
           out = model.generate(x,