katuni4ka commited on
Commit
6e6ca20
·
verified ·
1 Parent(s): 023f73f

Update modeling_arctic.py

Browse files
Files changed (1) hide show
  1. modeling_arctic.py +1 -1
modeling_arctic.py CHANGED
@@ -1766,7 +1766,7 @@ class ArcticForCausalLM(ArcticPreTrainedModel):
1766
  if isinstance(past_key_values, Cache):
1767
  cache_length = past_key_values.get_seq_length()
1768
  past_length = past_key_values.seen_tokens
1769
- max_cache_length = past_key_values.get_max_length()
1770
  else:
1771
  cache_length = past_length = past_key_values[0][0].shape[2]
1772
  max_cache_length = None
 
1766
  if isinstance(past_key_values, Cache):
1767
  cache_length = past_key_values.get_seq_length()
1768
  past_length = past_key_values.seen_tokens
1769
+ max_cache_length = past_key_values.get_max_length() if hasattr(past_key_values, "get_max_length") else past_key_values.get_max_cache_shape()
1770
  else:
1771
  cache_length = past_length = past_key_values[0][0].shape[2]
1772
  max_cache_length = None