Skip to content

GitLab

Commit 534fd185 authored by raojy's avatar raojy 💬
Browse files

Update run_Transformers_t2i_thin.py

parent fa537d0c
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
LLaDA2.0/run_Transformers_t2i_thin.py
View file @ 534fd185
...@@ -5,7 +5,7 @@ from decoder import decode_vq_tokens ...@@ -5,7 +5,7 @@ from decoder import decode_vq_tokens
model_path = "inclusionAI/LLaDA2.0-Uni" model_path = "inclusionAI/LLaDA2.0-Uni"
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
model_path, device_map="cuda", torch_dtype="bfloat16", trust_remote_code=True model_path, device_map="auto", torch_dtype="bfloat16", trust_remote_code=True
).eval() ).eval()
model.tokenizer = tokenizer model.tokenizer = tokenizer
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment