duzx16 commited on
Commit
f2191d0
1 Parent(s): b9c9fe6

Fix arange_cpu

Browse files
Files changed (1) hide show
  1. modeling_chatglm.py +1 -1
modeling_chatglm.py CHANGED
@@ -87,7 +87,7 @@ def split_tensor_along_last_dim(
87
  class RotaryEmbedding(nn.Module):
88
  def __init__(self, dim, original_impl=False, device=None, dtype=None):
89
  super().__init__()
90
- inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, device=device, dtype=dtype) / dim))
91
  self.register_buffer("inv_freq", inv_freq)
92
  self.dim = dim
93
  self.original_impl = original_impl
 
87
  class RotaryEmbedding(nn.Module):
88
  def __init__(self, dim, original_impl=False, device=None, dtype=None):
89
  super().__init__()
90
+ inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, device=device).to(dtype=dtype) / dim))
91
  self.register_buffer("inv_freq", inv_freq)
92
  self.dim = dim
93
  self.original_impl = original_impl