Update app.py

app.py

@@ -111,7 +111,7 @@ def load_model(model_selector):
     print('Model will use', dtype, 'precision...')
     print('=' * 70)
 
-    return model
+    return [model, ctx]
 
 #==================================================================================
 
@@ -207,7 +207,8 @@ def generate_music(prime,
     else:
         inputs = prime[-num_mem_tokens:]
 
-    model = model_state
+    model = model_state[0]
+    ctx = model_state[1]
         
     model.cuda()
     model.eval()
@@ -219,7 +220,6 @@ def generate_music(prime,
     inp = torch.LongTensor(inp).cuda()
     
     with ctx:
-      with torch.inference_mode():
         out = model.generate(inp,
                               num_gen_tokens,
                               #filter_logits_fn=top_p,