Bug: parameter used as a function in , causing NameError at inference

[kuishou68]

Bug Report

Description

In model/kronos.py, the function sample_from_logits (line 373–386) contains a logic bug on line 382: when sample_logits=False, the code attempts to call top_k as if it were a function, but top_k in that scope is an integer parameter, not a callable. This will raise a TypeError: 'int' object is not callable at runtime.

Location

File: model/kronos.py, function sample_from_logits, line 382

Buggy Code

def sample_from_logits(logits, temperature=1.0, top_k=None, top_p=None, sample_logits=True):
    logits = logits / temperature
    if top_k is not None or top_p is not None:
        if top_k > 0 or top_p < 1.0:
            logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)

    probs = F.softmax(logits, dim=-1)

    if not sample_logits:
        _, x = top_k(probs, k=1, dim=-1)   # ← BUG: top_k is an int, not a function
    else:
        x = torch.multinomial(probs, num_samples=1)

    return x

Root Cause

The intention is to take the argmax/greedy token when sample_logits=False. The correct PyTorch API is torch.topk(probs, k=1, dim=-1).

Fix

    if not sample_logits:
        _, x = torch.topk(probs, k=1, dim=-1)
    else:
        x = torch.multinomial(probs, num_samples=1)

Impact

Any call to sample_from_logits(..., sample_logits=False) will crash with TypeError: 'int' object is not callable. This path is used during greedy decoding in inference pipelines.

[kuishou68]

I've opened PR #232 to fix this issue.