kiranr commited on
Commit
da419fa
1 Parent(s): eae6b01

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +0 -5
handler.py CHANGED
@@ -2,10 +2,6 @@ import torch
2
  from typing import Dict, List, Any
3
  from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
4
 
5
- # check for GPU
6
- device = 0 if torch.cuda.is_available() else -1
7
-
8
-
9
  format_input = (
10
  "Below is an instruction that describes a task. "
11
  "Write a response that appropriately completes the request.\n\n"
@@ -27,7 +23,6 @@ class EndpointHandler:
27
  "text-generation",
28
  model=model,
29
  tokenizer=tokenizer,
30
- device=device,
31
  max_length=256,
32
  )
33
 
 
2
  from typing import Dict, List, Any
3
  from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
4
 
 
 
 
 
5
  format_input = (
6
  "Below is an instruction that describes a task. "
7
  "Write a response that appropriately completes the request.\n\n"
 
23
  "text-generation",
24
  model=model,
25
  tokenizer=tokenizer,
 
26
  max_length=256,
27
  )
28