fixed generate method

src/model.py

@@ -197,7 +197,10 @@ class GPT(nn.Module):
     def generate(self, input_ids, max_length=50,eos_token_id=None):
         generated_tokens = []
         current_ids = input_ids
-    
+        
+        # 🔥 Infer device from input_ids
+        device = input_ids.device
+
         for _ in range(max_length):
             # Forward pass to get logits
             logits = self.forward(current_ids)  # Shape: (batch_size, seq_len, vocab_size)