Update app.py

app.py

@@ -95,11 +95,13 @@ def GenerateMIDI(num_tok, idrums, iinstr):
         
         inp = torch.LongTensor([outy]).cpu()
 
-        out = model.module.generate(inp,
-                                    1,
-                                    temperature=0.9,
-                                    return_prime=False,
-                                    verbose=False)
+        with ctx:
+        
+            out = model.module.generate(inp,
+                                        1,
+                                        temperature=0.9,
+                                        return_prime=False,
+                                        verbose=False)
 
         out0 = out[0].tolist()
         outy.extend(out0)
@@ -177,12 +179,14 @@ if __name__ == "__main__":
     print('Loading model checkpoint...')
 
     model.load_state_dict(
-        torch.load('Allegro_Music_Transformer_Tiny_Trained_Model_80000_steps_0.9457_loss_0.7443_acc.pth',
+        torch.load('Ultimate_Drums_Transformer_Small_Trained_Model_8134_steps_0.3745_loss_0.8736_acc.pth',
                    map_location='cpu'))
     print('=' * 70)
 
     model.eval()
 
+    ctx = torch.amp.autocast(device_type='cpu', dtype=torch.bfloat16)
+
     print('Done!')
     print('=' * 70)