ahans30 commited on
Commit
ae02592
1 Parent(s): d389c0e

Added CPU support

Browse files
Files changed (1) hide show
  1. binoculars/detector.py +2 -1
binoculars/detector.py CHANGED
@@ -64,7 +64,8 @@ class Binoculars(object):
64
  def _get_logits(self, encodings: transformers.BatchEncoding) -> torch.Tensor:
65
  observer_logits = self.observer_model(**encodings.to(DEVICE_1)).logits
66
  performer_logits = self.performer_model(**encodings.to(DEVICE_2)).logits
67
- torch.cuda.synchronize()
 
68
  return observer_logits, performer_logits
69
 
70
  def compute_score(self, input_text: Union[list[str], str]) -> Union[float, list[float]]:
 
64
  def _get_logits(self, encodings: transformers.BatchEncoding) -> torch.Tensor:
65
  observer_logits = self.observer_model(**encodings.to(DEVICE_1)).logits
66
  performer_logits = self.performer_model(**encodings.to(DEVICE_2)).logits
67
+ if DEVICE_1 != "cpu":
68
+ torch.cuda.synchronize()
69
  return observer_logits, performer_logits
70
 
71
  def compute_score(self, input_text: Union[list[str], str]) -> Union[float, list[float]]: