zRzRzRzRzRzRzR hiyouga commited on
Commit
0891486
1 Parent(s): 6324425

fixes https://github.com/THUDM/GLM-4/issues/22 (#5)

Browse files

- Update modeling_chatglm.py (7122f9503a0d915971b96adad7efad64e3609177)


Co-authored-by: Yaowei Zheng <hiyouga@users.noreply.huggingface.co>

Files changed (1) hide show
  1. modeling_chatglm.py +1 -1
modeling_chatglm.py CHANGED
@@ -624,7 +624,7 @@ class GLMTransformer(torch.nn.Module):
624
  if len(presents) == 0:
625
  presents = kv_cache
626
  else:
627
- presents = torch.cat((presents, kv_cache), dim=0)
628
 
629
  if output_hidden_states:
630
  all_hidden_states = all_hidden_states + (hidden_states,)
 
624
  if len(presents) == 0:
625
  presents = kv_cache
626
  else:
627
+ presents = torch.cat((presents, kv_cache.to(presents.device)), dim=0)
628
 
629
  if output_hidden_states:
630
  all_hidden_states = all_hidden_states + (hidden_states,)