Chris4K commited on
Commit
8c9754a
·
verified ·
1 Parent(s): e0da0b4

Update services/model_service.py

Browse files
Files changed (1) hide show
  1. services/model_service.py +1 -0
services/model_service.py CHANGED
@@ -44,6 +44,7 @@ class ModelService:
44
  # Load model with the updated configuration
45
  self.model = AutoModelForCausalLM.from_pretrained(
46
  settings.MODEL_NAME,
 
47
  torch_dtype=torch.float16 if settings.DEVICE == "cuda" else torch.float32,
48
  device_map="auto" if settings.DEVICE == "cuda" else None,
49
  config=config
 
44
  # Load model with the updated configuration
45
  self.model = AutoModelForCausalLM.from_pretrained(
46
  settings.MODEL_NAME,
47
+ model_type = "llama",
48
  torch_dtype=torch.float16 if settings.DEVICE == "cuda" else torch.float32,
49
  device_map="auto" if settings.DEVICE == "cuda" else None,
50
  config=config