Update app.py

app.py

@@ -225,7 +225,7 @@ def generate_music(
                 print(
                     f'Section {i}: output length {input_ids.shape[-1]} exceeding context length {max_context}, now using the last {max_context} tokens.')
                 input_ids = input_ids[:, -(max_context):]
-            with torch.inference_mode(), torch.autocast(device_type='cuda', dtype=torch.float16):
+            with torch.inference_mode(), torch.autocast(device_type=device, dtype=torch.float16):
                 output_seq = model.generate(
                     input_ids=input_ids,
                     max_new_tokens=max_new_tokens,