diff --git a/modules/callbacks.py b/modules/callbacks.py index aa92f9cb..945b8c37 100644 --- a/modules/callbacks.py +++ b/modules/callbacks.py @@ -12,7 +12,7 @@ import modules.shared as shared # Copied from https://github.com/PygmalionAI/gradio-ui/ class _SentinelTokenStoppingCriteria(transformers.StoppingCriteria): - def __init__(self, sentinel_token_ids: list[torch.LongTensor], starting_idx: int): + def __init__(self, sentinel_token_ids: list, starting_idx: int): transformers.StoppingCriteria.__init__(self) self.sentinel_token_ids = sentinel_token_ids self.starting_idx = starting_idx