Consider refactoring CTC greedy decoding about nemo

[galv]

@artbataev this may be related to your observation of TDT being faster than CTC in some cases.

[titu1994]

If I recall correctly, max does work on GPU so I don't know why I forced it onto CPU. Maybe the slice? Either way, if it fixes the slowdown let's move the .CPU() to after max calculation.

It might have been that the logprobs can be super long for 12 hours so I wanted to preserve memory, in which case the cost of speed is not under consideration but I suppose we can make it a premature optimization and select whether to move the tensor to CPU before or after depending on length of the tensor.

[titu1994]

It is pretty common to run long audio CTC - NFA does it constantly. But on average, it is more common to run short audio batches

[galv]

It might have been that the logprobs can be super long for 12 hours so I wanted to preserve memory

The tensor pointed to by "prediction" already existed on GPU. It is just being copied to CPU in this code. We should not expect any memory increase from my suggestion. In fact, we will see a CPU memory decrease.

[titu1994]

That's true. Can you see if moving it onto CPU at the end before return is sufficient to remove this bottleneck then let's do it in your other pr