Update inference.py

inference.py

@@ -17,7 +17,7 @@ def format_chat(messages: List[Dict[str, str]], tokenizer) -> str:
     return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
 
 # Model loading function for SageMaker with tensor parallelism
-def model_fn(model_dir):
+def model_fn(model_dir,context=None):
     """
     Load the model and tokenizer from the model directory for inference.
     This version supports tensor parallelism across 4 GPUs.