Imran1
/

Qwen2.5-72B-Instruct-FP8

Model card Files Files and versions

Imran1 commited on Oct 10, 2024

Commit

54088e7

·

verified ·

1 Parent(s): 2744dc1

Update code/inference.py

Files changed (1) hide show

code/inference.py +1 -1

code/inference.py CHANGED Viewed

@@ -66,7 +66,7 @@ def model_fn(model_dir, context=None):
                 os.makedirs(offload_dir, exist_ok=True)
                 # Load and dispatch model across 4 GPUs using tensor parallelism
-                model = AutoModelForCausalLM.from_pretrained(model_dir, torch_dtype=torch.float16)
                 model = load_checkpoint_and_dispatch(
                     model,
                     model_dir,

                 os.makedirs(offload_dir, exist_ok=True)
                 # Load and dispatch model across 4 GPUs using tensor parallelism
+                model = AutoModelForCausalLM.from_pretrained(model_dir, torch_dtype="auto")
                 model = load_checkpoint_and_dispatch(
                     model,
                     model_dir,