pytorch - 💡(How to fix) Fix `torch.compile` raises InternalTorchDynamoError KeyError("class_type") when constructing a Tensor subclass [1 pull requests]

torch._dynamo.exc.InternalTorchDynamoError: KeyError: class_type

---

import os
import platform
import traceback

import torch


class MyTensor(torch.Tensor):
    pass


def f(x):
    y = MyTensor(x)
    return torch.abs(y)


def print_env():
    print("Python:", platform.python_version())
    print("Platform:", platform.platform())
    print("PyTorch:", torch.__version__)
    print("CUDA available:", torch.cuda.is_available())
    print("CUDA device count:", torch.cuda.device_count())
    print("CUDA_VISIBLE_DEVICES:", os.environ.get("CUDA_VISIBLE_DEVICES", ""))
    if torch.cuda.is_available():
        print("Current CUDA device:", torch.cuda.current_device())
        print("CUDA device name:", torch.cuda.get_device_name(0))


def main():
    print_env()

    if not torch.cuda.is_available():
        raise RuntimeError("This repro expects CUDA.")

    x = torch.randn(4, device="cuda")

    print("\nRunning eager...")
    eager_out = f(x)
    print("Eager succeeded.")
    print("Eager output type:", type(eager_out))
    print("Eager output:", eager_out)

    print("\nRunning compiled...")
    compiled_f = torch.compile(
        f,
        backend="inductor",
        fullgraph=True,
        dynamic=True,
    )

    try:
        compiled_out = compiled_f(x)
        print("Compiled succeeded.")
        print("Compiled output type:", type(compiled_out))
        print("Compiled output:", compiled_out)
        torch.testing.assert_close(
            torch.as_tensor(eager_out),
            torch.as_tensor(compiled_out),
        )
    except Exception:
        print("Compiled failed:")
        traceback.print_exc()
        raise


if __name__ == "__main__":
    main()

---

Running eager...
Eager succeeded.
Eager output type: <class '__main__.MyTensor'>

---

torch._dynamo.exc.InternalTorchDynamoError: KeyError: class_type

from user code:
   File "repro.py", line 16, in f
    y = MyTensor(x)

---

File "torch/_dynamo/variables/tensor.py", line 2585, in call_function
    var = TensorWithTFOverrideVariable.from_tensor_var(
File "torch/_dynamo/variables/torch_function.py", line 605, in from_tensor_var
    input_tensor_type = kwargs.pop("class_type")
KeyError: class_type

---

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