Joshua Lochner commited on
Commit
79b40d9
·
1 Parent(s): 3d1c770

Set device to cpu (-1) if device index is None

Browse files
Files changed (1) hide show
  1. src/classify.py +2 -0
src/classify.py CHANGED
@@ -6,6 +6,8 @@ import segment
6
  class SponsorBlockClassificationPipeline(TextClassificationPipeline):
7
  def __init__(self, model, tokenizer):
8
  device = next(model.parameters()).device.index
 
 
9
  super().__init__(model=model, tokenizer=tokenizer,
10
  return_all_scores=True, truncation=True, device=device)
11
 
 
6
  class SponsorBlockClassificationPipeline(TextClassificationPipeline):
7
  def __init__(self, model, tokenizer):
8
  device = next(model.parameters()).device.index
9
+ if device is None:
10
+ device = -1
11
  super().__init__(model=model, tokenizer=tokenizer,
12
  return_all_scores=True, truncation=True, device=device)
13