fix context size for llama2

Yisok1208 · Sep 28, 2023 · 069ac97 · 069ac97
1 parent 2928df3
commit 069ac97
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/lib/prune.py b/lib/prune.py
@@ -214,7 +214,7 @@ def tmp(_, inp, out):
 def prune_sparsegpt(args, model, tokenizer, dev, prune_n=0, prune_m=0):
     ## SparseGPT code available at: https://github.com/IST-DASLab/sparsegpt/tree/f5c25005a61f96a0933ca2f95705a963585aafaa
     print('Starting ...')
-    dataloader, _ = get_loaders("c4",nsamples=args.nsamples,seed=args.seed,seqlen=2048,tokenizer=tokenizer)
+    dataloader, _ = get_loaders("c4",nsamples=args.nsamples,seed=args.seed,seqlen=model.seqlen,tokenizer=tokenizer)
 
     use_cache = model.config.use_cache
     model.config.use_cache = False