pytorch - 💡(How to fix) Fix `torch.compile` fails on `nn.CTCLoss` with missing fake impl for `aten._use_cudnn_ctc_loss.Tensor`

import torch
import torch.nn as nn


class M(nn.Module):
    def __init__(self):
        super().__init__()
        self.loss = nn.CTCLoss(zero_infinity=True, reduction="sum")
        self.register_buffer("targets", torch.tensor([1, 2], dtype=torch.int32))
        self.register_buffer("input_lengths", torch.tensor([2, 2], dtype=torch.int32))
        self.register_buffer("target_lengths", torch.tensor([2, 0], dtype=torch.int32))

    def forward(self, x):
        log_probs = x.log_softmax(2)
        return self.loss(
            log_probs,
            self.targets,
            self.input_lengths,
            self.target_lengths,
        )


m = M().cuda().eval()
x = torch.rand(2, 2, 3, dtype=torch.double, device="cuda")

with torch.no_grad():
    eager = m(x)
    print("eager:", eager)

compiled = torch.compile(
    m,
    backend="inductor",
    fullgraph=True,
    dynamic=True,
)

with torch.no_grad():
    out = compiled(x)
    print("compiled:", out)

---

eager: tensor(4.6897, device='cuda:0', dtype=torch.float64)

---

NotImplementedError: aten::_use_cudnn_ctc_loss.Tensor: attempted to run this operator with Meta tensors,
but there was no fake impl or Meta kernel registered.

---

torch._dynamo.exc.Unsupported: Operator does not support running with fake tensors

Developer debug context: unsupported operator: aten._use_cudnn_ctc_loss.Tensor

from user code:
  return self.loss(log_probs, self.targets, self.input_lengths, self.target_lengths)

---

PyTorch version:  2.13.0a0+git059c270
Is debug build: True
CUDA used to build PyTorch: 12.6
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.4 LTS (x86_64)
GCC version: (Ubuntu 12.3.0-1ubuntu1~22.04) 12.3.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.10.16 (main, Dec 11 2024, 16:24:50) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-6.8.0-59-generic-x86_64-with-glibc2.35
Is CUDA available: True