Update app.py

app.py

@@ -19,7 +19,7 @@ encode = lambda s: [stoi[c] for c in s] # encoder: take a string, output a list
 decode = lambda l: ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string
 
 model = GPTLanguageModel(vocab_size)
-model.load_state_dict(torch.load('saved_model.pth'))
+model.load_state_dict(torch.load('saved_model.pth', map_location=cfg.device))
 m = model.to(cfg.device)
 
 def inference(input_context, count):
@@ -27,7 +27,7 @@ def inference(input_context, count):
     count = int(count)
     context = torch.tensor(encoded_text, dtype=torch.long, device=cfg.device)
      
-    print('--------------------context = ',context)
+    #print('--------------------context = ',context)
     out_text = decode(m.generate(context, max_new_tokens=count)[0].tolist())
     return out_text