Harry Coultas Blum commited on
Commit
17ad867
·
1 Parent(s): 76a4a55
Files changed (1) hide show
  1. vui/model.py +1 -1
vui/model.py CHANGED
@@ -38,7 +38,7 @@ class KVCache(nn.Module):
38
 
39
  k_out = self.k_cache
40
  v_out = self.v_cache
41
- input_pos = input_pos.to(k_out.dtype)
42
  k_out[:, :, input_pos] = k_val
43
  v_out[:, :, input_pos] = v_val
44
 
 
38
 
39
  k_out = self.k_cache
40
  v_out = self.v_cache
41
+ input_pos = input_pos.int()
42
  k_out[:, :, input_pos] = k_val
43
  v_out[:, :, input_pos] = v_val
44