added return_grad for all types of rnnt loss

durson commented 10 months ago

Grad tensors can be useful for the full (ordinary and pruned) rnnt losses.

pkufool commented 10 months ago

durson commented 10 months ago

I fixed #25 (comment) However I preserved S > 0 requirement since otherwise it is not compatible with torchaudio and unittest was not passing.

Summary of the changes:

Added return_grad for all rnnt losses.
Lifted T >= S requirement for regular case. That included making sure that shapes match not only for the whole tensor but for batch instances as well (testing boundaries if provided).
Added s_range adjustment as per D.Povey suggestion and added unit test to check for correctness.
Nitpicks - torch.Tensor -> Tensor, making bonudaries optional in appropriate cases.

pkufool commented 10 months ago

@durson I have added the github actions to this repo, could you sync the PR with the latest master, thanks!

durson commented 10 months ago

Merged latest master to fork and tests are passing.

pkufool commented 10 months ago

Thanks! merging! I would be very nice if you can also make a PR to https://github.com/k2-fsa/k2 (k2/python/k2/rnnt_loss.py).

k2-fsa / fast_rnnt