darkstar94 commited on
Commit
b92567e
1 Parent(s): d07861c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -3
app.py CHANGED
@@ -22,11 +22,12 @@ from transformers import AutoModel, AutoTokenizer
22
 
23
  # Argparser
24
  parser = argparse.ArgumentParser(description='demo')
25
- parser.add_argument('--device', type=str, default='mps', help='cuda or mps')
26
  parser.add_argument('--dtype', type=str, default='fp16', help='bf16 or fp16')
27
  args = parser.parse_args()
28
- device = args.device
29
- assert device in ['cuda', 'mps']
 
30
  if args.dtype == 'bf16':
31
  dtype = torch.bfloat16
32
  else:
 
22
 
23
  # Argparser
24
  parser = argparse.ArgumentParser(description='demo')
25
+ #parser.add_argument('--device', type=str, default='mps', help='cuda or mps')
26
  parser.add_argument('--dtype', type=str, default='fp16', help='bf16 or fp16')
27
  args = parser.parse_args()
28
+ #device = args.device
29
+ #assert device in ['cuda', 'mps']
30
+ device = 'cpu'
31
  if args.dtype == 'bf16':
32
  dtype = torch.bfloat16
33
  else: