Update app.py

app.py

@@ -37,7 +37,7 @@ CHECKPOINT_PATH=f'checkpoint-600'
 tokenizer = AutoTokenizer.from_pretrained("chatglm2-6b", trust_remote_code=True)
 config = AutoConfig.from_pretrained("chatglm2-6b", trust_remote_code=True, pre_seq_len=128)
 model = AutoModel.from_pretrained("chatglm2-6b", config=config, trust_remote_code=True)
-prefix_state_dict = torch.load(os.path.join(CHECKPOINT_PATH, "pytorch_model.bin"))
+prefix_state_dict = torch.load(os.path.join(CHECKPOINT_PATH, "pytorch_model.bin"), map_location=torch.device('cpu'))
 new_prefix_state_dict = {}
 for k, v in prefix_state_dict.items():
     if k.startswith("transformer.prefix_encoder."):