Harry Coultas Blum commited on
Commit
76a4a55
·
1 Parent(s): 8760fa5
Files changed (1) hide show
  1. vui/model.py +3 -2
vui/model.py CHANGED
@@ -38,8 +38,9 @@ class KVCache(nn.Module):
38
 
39
  k_out = self.k_cache
40
  v_out = self.v_cache
41
- k_out[:, :, input_pos] = k_val.to(k_out.dtype)
42
- v_out[:, :, input_pos] = v_val.to(v_val.dtype)
 
43
 
44
  return k_out, v_out
45
 
 
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
 
45
  return k_out, v_out
46