rwitz commited on
Commit
b6896e7
1 Parent(s): 2d49749

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +1 -1
handler.py CHANGED
@@ -6,7 +6,7 @@ sleep_time = int(os.environ.get('SLEEP_TIME', 3))
6
  # Use a pipeline as a high-level helper
7
  from transformers import pipeline
8
 
9
- pipe = pipeline("text-generation", model="rwitz/go-bruins-v2",device=0,torch_dtype=torch.bfloat16)
10
  ## load your model(s) into vram here
11
 
12
  def handler(event):
 
6
  # Use a pipeline as a high-level helper
7
  from transformers import pipeline
8
 
9
+ pipe = pipeline("text-generation", model="rwitz/go-bruins-v2",device=0,torch_dtype=torch.bfloat16,cache_dir="/workspace")
10
  ## load your model(s) into vram here
11
 
12
  def handler(event):