Update code/inference.py

code/inference.py

@@ -2,7 +2,7 @@ import os
 import json
 import torch
 import torch.nn as nn
-from transformers import AutoModel, AutoTokenizer, AutoConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
 import logging
 
 logger = logging.getLogger(__name__)
@@ -43,11 +43,10 @@ def model_fn(model_dir, context=None):
         
         # Load config and specify it's a Phi3Config
         config = AutoConfig.from_pretrained(model_id, 
-                                          num_labels=2,
                                           trust_remote_code=True)
         
-        # Load base model
-        base_model = AutoModel.from_pretrained(
+        # Load base model using AutoModelForCausalLM
+        base_model = AutoModelForCausalLM.from_pretrained(
             model_id,
             config=config,
             torch_dtype=torch.bfloat16 if device.type == 'cuda' else torch.float32,