BK-Lee commited on
Commit
891b91b
1 Parent(s): 2053e3b
Files changed (1) hide show
  1. app.py +7 -1
app.py CHANGED
@@ -14,9 +14,15 @@ from transformers import TextIteratorStreamer
14
  from torchvision.transforms.functional import pil_to_tensor
15
 
16
  # loading meteor model
17
- mmamba = load_mmamba('BK-Lee/Meteor-Mamba').cuda()
18
  meteor, tok_meteor = load_meteor('BK-Lee/Meteor-MLM', bits=4)
19
 
 
 
 
 
 
 
20
  # device
21
  device = torch.cuda.current_device()
22
 
 
14
  from torchvision.transforms.functional import pil_to_tensor
15
 
16
  # loading meteor model
17
+ mmamba = load_mmamba('BK-Lee/Meteor-Mamba').to('cuda')
18
  meteor, tok_meteor = load_meteor('BK-Lee/Meteor-MLM', bits=4)
19
 
20
+ # param
21
+ for param in mmamba.parameters():
22
+ param = param.to('cuda')
23
+ for param in meteor.parameters():
24
+ param = param.to('cuda')
25
+
26
  # device
27
  device = torch.cuda.current_device()
28