I’ve tried  the following:
def ctc_alignment_targets(logits, targets, input_lengths, target_lengths, blank, dim = 1):
    with torch.enable_grad():
        logits = logits.detach().requires_grad_()
        log_probs = F.log_softmax(logits, dim = dim)
        ctc_loss = F.ctc_loss(log_probs.permute(2, 0, 1), targets, input_lengths, target_lengths, blank, reduction = 'sum')
        ctc_grad, = torch.autograd.grad(ctc_loss, (logits,))
        temporal_mask = (torch.arange(log_probs.shape[-1], device = input_lengths.device, dtype = input_lengths.dtype).unsqueeze(0) < input_lengths.unsqueeze(1))[:, None, :]
        alignment_targets = log_probs.exp() * temporal_mask - ctc_grad
        return log_probs, ctc_loss, alignment_targets

log_probs, ctc_loss, alignment_targets = ctc_alignment_targets(logits, ...)
ctc_loss_via_mle = (-alignment_targets * log_probs).sum()
print(ctc_loss, ctc_loss_via_mle)

Something it not quite right, the two quantities don’t match. Maybe did sth wrong, but the autograd.grad approach worked (even without a custom autograd function, though not sure if it’s same speed as the custom autograd function approach you suggested)