Skip to content

Commit 0a3bd13

Browse files
minor fix for model's device
1 parent abb2f98 commit 0a3bd13

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

detectors/huggingface/detector.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -286,7 +286,7 @@ def close(self) -> None:
286286
"""Clean up model and tokenizer resources."""
287287

288288
if self.model:
289-
if hasattr(self.model, 'to') and self.cuda_device.type == "cuda":
289+
if hasattr(self.model, 'to') and hasattr(self.model, 'device') and self.model.device.type != "cpu":
290290
self.model = self.model.to(torch.device("cpu"))
291291
self.model = None
292292

0 commit comments

Comments
 (0)