Update app.py

app.py

@@ -57,8 +57,6 @@ def generate_interactive(
 ):
     inputs = tokenizer([prompt], padding=True, return_tensors='pt')
     input_length = len(inputs['input_ids'][0])
-    for k, v in inputs.items():
-        inputs[k] = v.cuda()
     input_ids = inputs['input_ids']
     _, input_ids_seq_length = input_ids.shape[0], input_ids.shape[-1]
     if generation_config is None:
@@ -184,7 +182,7 @@ def load_model():
     model_dir = 'internlm/internlm2-chat-1_8b'
     model = (AutoModelForCausalLM.from_pretrained(
         model_dir,
-        trust_remote_code=True).to(torch.bfloat16).cuda())
+        trust_remote_code=True).to(torch.bfloat16))
     tokenizer = AutoTokenizer.from_pretrained(
         model_dir,
         trust_remote_code=True)
@@ -232,7 +230,6 @@ def combine_history(prompt):
 
 
 def main():
-    # torch.cuda.empty_cache()
     print('load model begin.')
     model, tokenizer = load_model()
     print('load model end.')
@@ -278,7 +275,7 @@ def main():
             'role': 'robot',
             'content': cur_response,  # pylint: disable=undefined-loop-variable
         })
-        torch.cuda.empty_cache()
+
 
 
 if __name__ == '__main__':