jameshuntercarter commited on
Commit
37eb59e
1 Parent(s): b5dba8a

Upload predict.py

Browse files
Files changed (1) hide show
  1. predict.py +2 -1
predict.py CHANGED
@@ -20,7 +20,7 @@ class Predictor(BasePredictor):
20
  """Run a single prediction on the model"""
21
  # SETUP
22
  large_quant_model = False # Use the larger pretrained model
23
- device = 'cuda' # 'cuda', 'cpu', 'cuda:0', 0, -1, torch.device('cuda')
24
 
25
  import numpy as np
26
  import torch
@@ -35,6 +35,7 @@ class Predictor(BasePredictor):
35
  'quantifier_hubert_base_ls960_14.pth', 'tokenizer.pth')
36
 
37
  print('Loading HuBERT...')
 
38
  hubert_model = CustomHubert(
39
  HuBERTManager.make_sure_hubert_installed(), device=device)
40
  print('Loading Quantizer...')
 
20
  """Run a single prediction on the model"""
21
  # SETUP
22
  large_quant_model = False # Use the larger pretrained model
23
+ device = torch.device('cuda') # 'cuda', 'cpu', 'cuda:0', 0, -1, torch.device('cuda')
24
 
25
  import numpy as np
26
  import torch
 
35
  'quantifier_hubert_base_ls960_14.pth', 'tokenizer.pth')
36
 
37
  print('Loading HuBERT...')
38
+
39
  hubert_model = CustomHubert(
40
  HuBERTManager.make_sure_hubert_installed(), device=device)
41
  print('Loading Quantizer...')