Update inference/infer.py

inference/infer.py

@@ -139,6 +139,9 @@ repetition_penalty = 1.2
 # special tokens
 start_of_segment = mmtokenizer.tokenize('[start_of_segment]')
 end_of_segment = mmtokenizer.tokenize('[end_of_segment]')
+
+raw_output = None
+
 # Format text prompt
 run_n_segments = min(args.run_n_segments+1, len(lyrics))
 
@@ -197,6 +200,7 @@ for i, p in enumerate(tqdm(prompt_texts[:run_n_segments])):
         raw_output = torch.cat([raw_output, prompt_ids, output_seq[:, input_ids.shape[-1]:]], dim=1)
     else:
         raw_output = output_seq
+    print(len(raw_output))
 
 # save raw output and check sanity
 ids = raw_output[0].cpu().numpy()