In sample_from_logits(), the branch for sample_logits=False incorrectly
called 'top_k' as if it were a function (top_k(probs, k=1, dim=-1)),
but 'top_k' is an integer parameter in this scope. This raises:
TypeError: 'int' object is not callable
Fix: replace with torch.topk(probs, k=1, dim=-1) which is the correct
PyTorch API for greedy (argmax) token selection.
Also fix a docstring typo in decode_s2(): 'torch.torch.Tensor' -> 'torch.Tensor'
Closes#231
Without releasing cached GPU memory, usage will keep growing during autoregressive prediction, leading to significant memory increase or OOM. Calling torch.cuda.empty_cache() prevents this accumulation.
cocoon