foxxy-hm commited on
Commit
b2f8848
·
1 Parent(s): 4628cd1

Update src/models/pairwise_model.py

Browse files
Files changed (1) hide show
  1. src/models/pairwise_model.py +1 -1
src/models/pairwise_model.py CHANGED
@@ -14,7 +14,7 @@ pad_token_id = tokenizer.pad_token_id
14
 
15
 
16
  class PairwiseModel(nn.Module):
17
- def __init__(self, model_name, max_length=384, batch_size=16, device="cuda:0"):
18
  super(PairwiseModel, self).__init__()
19
  self.max_length = max_length
20
  self.batch_size = batch_size
 
14
 
15
 
16
  class PairwiseModel(nn.Module):
17
+ def __init__(self, model_name, max_length=384, batch_size=16, device="cpu"):
18
  super(PairwiseModel, self).__init__()
19
  self.max_length = max_length
20
  self.batch_size = batch_size