fixes https://github.com/THUDM/GLM-4/issues/22

#1
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,)