ikeno-ada commited on
Commit
1bac627
1 Parent(s): e3c5e74

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +1 -1
handler.py CHANGED
@@ -7,7 +7,7 @@ import torch
7
  class EndpointHandler():
8
  def __init__(self, path=""):
9
  # load the optimized model
10
- self.model = M2M100ForConditionalGeneration.from_pretrained(path, device_map= "cuda",torch_dtype=torch.bfloat16)
11
  self.tokenizer = M2M100Tokenizer.from_pretrained(path)
12
 
13
  def __call__(self, data: Dict[str,str]) -> Dict[str, str]:
 
7
  class EndpointHandler():
8
  def __init__(self, path=""):
9
  # load the optimized model
10
+ self.model = M2M100ForConditionalGeneration.from_pretrained(path,torch_dtype=torch.bfloat16)
11
  self.tokenizer = M2M100Tokenizer.from_pretrained(path)
12
 
13
  def __call__(self, data: Dict[str,str]) -> Dict[str, str]: