anirudh-sub commited on
Commit
de2fd00
·
1 Parent(s): b45ed0e

Added torch.dtype

Browse files
Files changed (1) hide show
  1. app.py +1 -0
app.py CHANGED
@@ -7,6 +7,7 @@ tokenizer = AutoTokenizer.from_pretrained(model)
7
  pipeline = transformers.pipeline(
8
  "text-generation",
9
  model=model,
 
10
  device_map="auto",
11
  )
12
 
 
7
  pipeline = transformers.pipeline(
8
  "text-generation",
9
  model=model,
10
+ torch_dtype=torch.float16,
11
  device_map="auto",
12
  )
13