jeffreykthomas commited on
Commit
2927230
·
1 Parent(s): 96f6ccb

Changed dtype

Browse files
Files changed (1) hide show
  1. handler.py +1 -1
handler.py CHANGED
@@ -3,7 +3,7 @@ from transformers import GenerationConfig, pipeline
3
  import torch
4
  from typing import Any, Dict
5
 
6
- dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
7
 
8
 
9
  class EndpointHandler:
 
3
  import torch
4
  from typing import Any, Dict
5
 
6
+ dtype = torch.float16
7
 
8
 
9
  class EndpointHandler: