trying to fix float error

app.py

@@ -39,7 +39,18 @@ def load_model(repo_id, local_file):
     model.eval()
     return model
 
-def generate(model, prompt, max_tokens, temperature=0.7):
+# def generate(model, prompt, max_tokens, temperature=0.7):
+#     for _ in range(max_tokens):
+#         prompt = prompt[:, :config.context_len]
+#         logits = model(prompt)
+#         logits = logits[:, -1, :] / temperature
+#         logit_probs = nn.functional.softmax(logits, dim=-1)
+#         next_prompt = torch.multinomial(logit_probs, num_samples=1)
+#         prompt = torch.cat((prompt, next_prompt), dim=1)
+#     return prompt
+
+def generate(model, input_ids, max_tokens, temperature=0.7):
+    prompt = input_ids
     for _ in range(max_tokens):
         prompt = prompt[:, :config.context_len]
         logits = model(prompt)