buxtcodes commited on
Commit
2fc97fa
·
verified ·
1 Parent(s): 9619ab5
Files changed (1) hide show
  1. custom_generate/generate.py +1 -0
custom_generate/generate.py CHANGED
@@ -2,6 +2,7 @@ import torch
2
  from typing import Any, Dict, List, Optional, Tuple
3
 
4
  from transformers import Cache, GenerationConfig
 
5
 
6
 
7
  UNSUPPORTED_GENERATION_ARGS = [
 
2
  from typing import Any, Dict, List, Optional, Tuple
3
 
4
  from transformers import Cache, GenerationConfig
5
+ from transformers.cache_utils import CacheLayer
6
 
7
 
8
  UNSUPPORTED_GENERATION_ARGS = [