Spaces:

m-ric
/

attribution_visualizer

Runtime error

m-ric HF staff commited on Jan 15

Commit

d4bacf6

1 Parent(s): c7c30cc

Set device map auto

Files changed (1) hide show

app.py CHANGED Viewed

@@ -10,7 +10,7 @@ from lxt.models.llama import LlamaForCausalLM, attnlrp
 from lxt.utils import clean_tokens
 # Load model and tokenizer
-model = LlamaForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0", torch_dtype=torch.bfloat16, device_map="cuda")
 tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
 # Apply AttnLRP rules

 from lxt.utils import clean_tokens
 # Load model and tokenizer
+model = LlamaForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0", torch_dtype=torch.bfloat16, device_map="auto")
 tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
 # Apply AttnLRP rules