Repository: pytorch/TensorRT

Issue:
🐛 [Bug] Incorrect judgment of model type
##  Bug Description

When I provide a TorchScript model, Torch-TensorRT claims it is an nn.Module model

```
INFO: [Torch-TensorRT] - ir was set to default, using TorchScript as ir
INFO: [Torch-TensorRT] - Module was provided as a torch.nn.Module, trying to script the module with torch.jit.script. In the event of a failure please preconvert your module to TorchScript
```

<!-- A clear and concise description of what the bug is. -->

## To Reproduce

Just run:

```
import torch
import torch.nn as nn
import torch_tensorrt as tt

class Model(nn.Module):
    def __init__(self):
        super(Model,self).__init__()
        self.conv=nn.Conv2d(3,3,3)

    def forward(self,x1):
        out=self.conv(x1)
        return out

a=Model().cuda().eval()
b=torch.jit.trace(a,torch.ones([1,3,20,20]).cuda())

compile_settings = {}
compile_settings["inputs"] = [tt.Input(shape = [1,3,20,20])]
tt.logging.set_reportable_log_level(tt.logging.Level.Graph)

tt.compile(b,**compile_settings)
```

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

Model type is correctly judged.

<!-- A clear and concise description of what you expected to happen. -->

## Environment

> Build information about Torch-TensorRT can be found by turning on debug messages

 - Torch-TensorRT Version : 1.0.0
 - PyTorch Version: 1.10.0
 - CPU Architecture: Intel(R) Xeon(R) Platinum 8352Y CPU @ 2.20GHz
 - OS (e.g., Linux): CentOS 7
 - How you installed PyTorch (`conda`, `pip`, `libtorch`, source): pip
 - Build command you used (if compiling from source):
 - Are you using local sources or building from archives:
 - Python version: 3.6.8
 - CUDA version: 11.4
 - GPU models and configuration: A30
 - Any other relevant information:

## Additional context

I have fixed this bug locally. After confirming the existence of this bug, I will open a PR to fix it.

<!-- Add any other context about the problem here. -->



Discussion:

Resolving patch:
diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py
--- a/py/torch_tensorrt/_compile.py
+++ b/py/torch_tensorrt/_compile.py
@@ -14,11 +14,28 @@ class _IRType(Enum):
     fx = 1
 
 
-def _module_ir(module: Any, ir: str) -> _IRType.ts:
-    # Possible module types
-    module_is_tsable = any(
-        isinstance(module, t) for t in [torch.nn.Module, torch.jit.ScriptModule, torch.jit.ScriptFunction])
-    module_is_fxable = any(isinstance(module, t) for t in [torch.nn.Module, torch.fx.GraphModule])
+class _ModuleType(Enum):
+    """Enum to set the minimum required logging level to print a message to stdout
+    """
+    nn = 0
+    ts = 1
+    fx = 2
+
+
+def _parse_module_type(module: Any) -> _ModuleType:
+    if any(isinstance(module, t) for t in [torch.jit.ScriptModule, torch.jit.ScriptFunction]):
+        return _ModuleType.ts
+    elif isinstance(module, torch.fx.GraphModule):
+        return _ModuleType.fx
+    elif isinstance(module, torch.nn.Module):
+        return _ModuleType.nn
+    else:
+        raise RuntimeError("Module is an unknown format")
+
+
+def _get_target_ir(module_type: _ModuleType, ir: str) -> _IRType:
+    module_is_tsable = any([module_type == t for t in [_ModuleType.nn, _ModuleType.ts]])
+    module_is_fxable = any([module_type == t for t in [_ModuleType.nn, _ModuleType.fx]])
 
     ir_targets_torchscript = any([ir == opt for opt in ["torchscript", "ts"]])
     ir_targets_fx = ir == "fx"
@@ -26,7 +43,7 @@ def _module_ir(module: Any, ir: str) -> _IRType.ts:
     if module_is_tsable and ir_targets_torchscript:
         return _IRType.ts
     elif module_is_fxable and ir_targets_fx:
-        if isinstance(module, torch.fx.GraphModule):
+        if module_type == _ModuleType.fx:
             raise ValueError("Was given a torch.fx.GraphModule, fx is not currently supported by Torch-TensorRT")
         elif ir_targets_fx:
             raise ValueError("Preferred ir was set to \"fx\" which is currently not supported by Torch-TensorRT")
@@ -85,10 +102,11 @@ def compile(module: Any, ir="default", inputs=[], enabled_precisions=set([_enums
     Returns:
         torch.nn.Module: Compiled Module, when run it will execute via TensorRT
     """
-    target_ir = _module_ir(module, ir)
+    module_type = _parse_module_type(module)
+    target_ir = _get_target_ir(module_type, ir)
     if target_ir == _IRType.ts:
         ts_mod = module
-        if isinstance(module, torch.nn.Module):
+        if module_type == _ModuleType.nn:
             logging.log(
                 logging.Level.Info,
                 "Module was provided as a torch.nn.Module, trying to script the module with torch.jit.script. In the event of a failure please preconvert your module to TorchScript"
@@ -134,14 +152,14 @@ def convert_method_to_trt_engine(module: Any,
         enabled_precision (Set(Union(torch.dtype, torch_tensorrt.dtype))): The set of datatypes that TensorRT can use when selecting kernels
         ir (str): The requested strategy to compile. (Options: default - Let Torch-TensorRT decide, ts - TorchScript with scripting path)
         **kwargs: Additional settings for the specific requested strategy (See submodules for more info)
-
     Returns:
         bytes: Serialized TensorRT engine, can either be saved to a file or deserialized via TensorRT APIs
     """
-    target_ir = _module_ir(module, ir)
+    module_type = _parse_module_type(module)
+    target_ir = _get_target_ir(module_type, ir)
     if target_ir == _IRType.ts:
         ts_mod = module
-        if isinstance(module, torch.nn.Module):
+        if module_type == _ModuleType.nn:
             logging.log(
                 logging.Level.Info,
                 "Module was provided as a torch.nn.Module, trying to script the module with torch.jit.script. In the event of a failure please preconvert your module to TorchScript"
@@ -155,4 +173,4 @@ def convert_method_to_trt_engine(module: Any,
     elif target_ir == _IRType.fx:
         raise RuntimeError("fx is currently not supported")
     else:
-        raise RuntimeError("Module is an unknown format or the ir requested is unknown")
+        raise RuntimeError("Module is an unknown format or the ir requested is unknown")
\ No newline at end of file


Test patch:
diff --git a/tests/py/test_api.py b/tests/py/test_api.py
--- a/tests/py/test_api.py
+++ b/tests/py/test_api.py
@@ -506,6 +506,18 @@ def test_dynamic_shape(self):
         self.assertTrue(self._verify_correctness(i, target))
 
 
+class TestModule(unittest.TestCase):
+
+    def test_module_type(self):
+        nn_module = models.alexnet(pretrained=True).eval().to("cuda")
+        ts_module = torch.jit.trace(nn_module, torch.ones([1, 3, 224, 224]).to("cuda"))
+        fx_module = torch.fx.symbolic_trace(nn_module)
+
+        self.assertEqual(torchtrt._compile._parse_module_type(nn_module), torchtrt._compile._ModuleType.nn)
+        self.assertEqual(torchtrt._compile._parse_module_type(ts_module), torchtrt._compile._ModuleType.ts)
+        self.assertEqual(torchtrt._compile._parse_module_type(fx_module), torchtrt._compile._ModuleType.fx)
+
+
 def test_suite():
     suite = unittest.TestSuite()
     suite.addTest(unittest.makeSuite(TestLoggingAPIs))
@@ -527,6 +539,7 @@ def test_suite():
     suite.addTest(unittest.makeSuite(TestCheckMethodOpSupport))
     suite.addTest(unittest.makeSuite(TestDevice))
     suite.addTest(unittest.makeSuite(TestInput))
+    suite.addTest(unittest.makeSuite(TestModule))
 
     return suite