Fix top k generation for k != 0
Seems like the shapes didn't line up for the comparison. Logits are `(batch_size, values)`. The minima had shape `(batchsize)` and couldn't be directly compared
Thanks @CatalinVoss!