🐛 Describe the bug

from enum import Enum
import torch

class Paint(torch.nn.Module):
    class Color(Enum):
        RED = "1"
        GREEN = "2"

@torch.jit.script
def enum_fn(x: Paint.Color, y: Paint.Color) -> bool:
    if x == Paint.Color.RED:
        return True
    return x == y

The code would work fine if Paint does not inherit nn.Module but here is the output if it does:

Traceback (most recent call last):
  File "torch_script_loss.py", line 16, in <module>
    def enum_fn(x: Paint.Color, y: Paint.Color) -> bool:
  File "/opt/conda/lib/python3.8/site-packages/torch/jit/_script.py", line 1322, in script
    fn = torch._C._jit_script_compile(
RuntimeError: 
attribute lookup is not defined on python value of type 'type':
  File "torch_script_loss.py", line 17
@torch.jit.script
def enum_fn(x: Paint.Color, y: Paint.Color) -> bool:
    if x == Paint.Color.RED:
            ~~~~~~~~~~~ <--- HERE
        return True
    return x == y

Versions

PyTorch version: 1.12.0a0+2c916ef
Is debug build: False
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04) 9.4.0
Clang version: Could not collect
CMake version: version 3.22.3
Libc version: glibc-2.31

Python version: 3.8.12 | packaged by conda-forge | (default, Jan 30 2022, 23:42:07) [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.13.0-30-generic-x86_64-with-glibc2.10
Is CUDA available: True
CUDA runtime version: 11.6.112

Nvidia driver version: 470.103.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.3.3
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.3.3
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.3.3
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.3.3
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.3.3
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.3.3
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.3.3
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy==0.942
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.22.3
[pip3] pytorch-ignite==0.4.8
[pip3] pytorch-quantization==2.1.2
[pip3] torch==1.12.0a0+2c916ef
[pip3] torch-tensorrt==1.1.0a0
[pip3] torchinfo==1.6.5
[pip3] torchtext==0.12.0a0
[pip3] torchvision==0.13.0a0
[conda] magma-cuda110 2.5.2 5 local
[conda] mkl 2019.5 281 conda-forge
[conda] mkl-include 2019.5 281 conda-forge
[conda] numpy 1.22.3 py38h05e7239_0 conda-forge
[conda] pytorch-ignite 0.4.8 pypi_0 pypi
[conda] pytorch-quantization 2.1.2 pypi_0 pypi
[conda] torch 1.12.0a0+2c916ef pypi_0 pypi
[conda] torch-tensorrt 1.1.0a0 pypi_0 pypi
[conda] torchinfo 1.6.5 pypi_0 pypi
[conda] torchtext 0.12.0a0 pypi_0 pypi
[conda] torchvision 0.13.0a0 pypi_0 pypi