Spaces:
Running
on
Zero
Running
on
Zero
Added CPU support
Browse files- 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 |
-
|
|
|
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]]:
|