Update code/inference.py

code/inference.py

@@ -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 = AutoModelForCausalLM.from_pretrained(model_dir, torch_dtype="auto")
                 model = load_checkpoint_and_dispatch(
                     model,
                     model_dir,