rlanner-echocap commited on
Commit
f9a0290
1 Parent(s): 09d8478

Update handler.py

Browse files

Removing load in 8-bit flag

Files changed (1) hide show
  1. handler.py +1 -1
handler.py CHANGED
@@ -11,7 +11,7 @@ class EndpointHandler:
11
  def __init__(self, path=""):
12
  # load the model
13
  tokenizer = AutoTokenizer.from_pretrained(path)
14
- model = AutoModelForCausalLM.from_pretrained(path, load_in_8bit=True, device_map="auto", torch_dtype=dtype, trust_remote_code=True)
15
 
16
  # create inference pipeline
17
  self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, device='cuda:0')
 
11
  def __init__(self, path=""):
12
  # load the model
13
  tokenizer = AutoTokenizer.from_pretrained(path)
14
+ model = AutoModelForCausalLM.from_pretrained(path, device_map="auto", torch_dtype=dtype, trust_remote_code=True)
15
 
16
  # create inference pipeline
17
  self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, device='cuda:0')