Fix inference

app.py

@@ -21,7 +21,7 @@ dropout = 0.4
 tokenizer = tiktoken.encoding_for_model('gpt2')
 model_path = hf_hub_download('cbspace/gpt', 'model.safetensors')
 state_dict = load_file(model_path)
-model = GPTModel(n_layers, n_heads, embed_dim, ffn_dim, n_vocab, max_seq_len, dropout)
+model = GPTModel(device, n_layers, n_heads, embed_dim, ffn_dim, n_vocab, max_seq_len, dropout)
 model.load_state_dict(state_dict, strict=False)
 model.eval()
 

model.py

@@ -24,9 +24,10 @@ class TransformerBlock(nn.Module):
 
 
 class GPTModel(nn.Module):
-    def __init__(self, n_layers, n_heads, embed_dim, ffn_dim, n_vocab, max_seq_len, dropout):
+    def __init__(self, device, n_layers, n_heads, embed_dim, ffn_dim, n_vocab, max_seq_len, dropout):
         super().__init__()
         self.max_seq_len = max_seq_len
+        self.device = device
         self.embedding = nn.Embedding(n_vocab, embed_dim)
         nn.init.normal_(self.embedding.weight, mean=0.0, std=0.02)
         self.positional_embedding = nn.Embedding(max_seq_len, embed_dim)
@@ -55,7 +56,7 @@ class GPTModel(nn.Module):
         context_list = [i for i in input_ctx]
         with torch.no_grad():
             while len(context_list) < max_length:
-                context = torch.tensor(context_list, dtype=torch.long, device=device).unsqueeze(0)
+                context = torch.tensor(context_list, dtype=torch.long, device=self.device).unsqueeze(0)
                 logits = self(context)[0,-1,:]
                 
                 if greedy: