Skip to content

Adds batching and CTCLoss functionality to RNNProcessor#16

Open
Sam-NT wants to merge 2 commits into
devfrom
nt/rnn-batch
Open

Adds batching and CTCLoss functionality to RNNProcessor#16
Sam-NT wants to merge 2 commits into
devfrom
nt/rnn-batch

Conversation

@Sam-NT

@Sam-NT Sam-NT commented Jul 7, 2026

Copy link
Copy Markdown

This PR adds support for RNNProcessor.partial_fit to fit on a batch (multiple samples of sequences of features) instead of one sample at a time. It further adds support for using torch.nn.CTCLoss as a customized loss function (similar to torch.nn.CrossEntropyLoss), due to its atypical function arguments.

I roughly designed the data structure containing sequence and target lengths to match the existing method for passing targets and data in one AxisArray, but it's not the cleanest. Happy to work through a better implementation in this PR if desired.

Also worth noting that torch.nn.CTCLoss is not implemented in MPS in the latest release, but is in active development. The latest nightly worked for me.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant