improve web ui

Former-commit-id: 5c0148c018b12b52bc5748acfd6ad43836f2edb5
This commit is contained in:
hiyouga
2024-01-10 12:37:45 +08:00
parent 6fa6d4532e
commit 64246d42d2
4 changed files with 26 additions and 14 deletions

View File

@@ -79,6 +79,10 @@ def get_current_device() -> torch.device:
return torch.device(device)
def get_device_count() -> int:
return torch.cuda.device_count()
def get_logits_processor() -> "LogitsProcessorList":
r"""
Gets logits processor that removes NaN and Inf logits.