Function next_token

lit_gpt/generate/base.py:59–64

(
    model: GPT, input_pos: torch.Tensor, x: torch.Tensor, **kwargs: Any
)

def next_token(
    model: GPT, input_pos: torch.Tensor, x: torch.Tensor, **kwargs: Any
) -> torch.Tensor:
    logits = model(x, input_pos)
    next = sample(logits, **kwargs)
    return next.to(dtype=x.dtype)


@torch.inference_mode()