Tolerate torch.jit.script call to Enum classes

test/jit/test_enum.py

@@ -348,3 +348,13 @@ def iterate_enum(x: Color):
         # PURPLE always appears last because we follow Python's Enum definition order.
         self.assertEqual(scripted(Color.RED), [Color.GREEN.value, Color.BLUE.value])
         self.assertEqual(scripted(Color.GREEN), [Color.RED.value, Color.BLUE.value])
+
+    # Tests that explicitly and/or repeatedly scripting an Enum class is permitted.
+    def test_enum_explicit_script(self):
+
+        @torch.jit.script
+        class Color(Enum):
+            RED = 1
+            GREEN = 2
+
+        torch.jit.script(Color)

torch/jit/_script.py

@@ -8,6 +8,7 @@
 """
 import functools
 import collections
+import enum
 import inspect
 import copy
 import pickle
@@ -949,6 +950,11 @@ def forward(self, input):
                 " pass an instance instead".format(obj)
             )
 
+        # Enums are automatically usable in TorchScript, explicitly scripting
+        # is not necessary, but not harmful either.
+        if issubclass(obj, enum.Enum):
+            return obj
+
         if not _is_new_style_class(obj):
             raise RuntimeError(
                 "TorchScript classes must be new-style classes. "