From d43a79b7bfb3be0b43c7afebf4cfff306ef92185 Mon Sep 17 00:00:00 2001 From: Casper Date: Sun, 17 Mar 2024 19:52:56 +0100 Subject: [PATCH] device_map auto --- examples/mistral/mixtral_fused.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/mistral/mixtral_fused.py b/examples/mistral/mixtral_fused.py index 5aaf9652b..5e72e2266 100644 --- a/examples/mistral/mixtral_fused.py +++ b/examples/mistral/mixtral_fused.py @@ -5,7 +5,7 @@ from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock model_path = "mistralai/Mixtral-8x7B-Instruct-v0.1" # Load model -model = AutoModelForCausalLM.from_pretrained(model_path) +model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto") for name, module in model.named_modules(): if isinstance(module, MixtralSparseMoeBlock):