oleksandrfluxon commited on
Commit
4d0044f
1 Parent(s): f3b8e58

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +1 -0
handler.py CHANGED
@@ -22,6 +22,7 @@ class EndpointHandler:
22
  with torch.autocast('cuda'):
23
  # preprocess
24
  inputs = self.tokenizer(inputs, return_tensors="pt").to('cuda')
 
25
 
26
  # pass inputs with all kwargs in data
27
  if parameters is not None:
 
22
  with torch.autocast('cuda'):
23
  # preprocess
24
  inputs = self.tokenizer(inputs, return_tensors="pt").to('cuda')
25
+ self.model.to('cuda')
26
 
27
  # pass inputs with all kwargs in data
28
  if parameters is not None: