We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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. -->


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `py/torch_tensorrt/_compile.py`
Content:
```
1 from typing import List, Dict, Any
2 from torch_tensorrt import _enums
3 import torch_tensorrt.ts
4 from torch_tensorrt import logging
5 import torch
6 from torch import fx
7 from enum import Enum
8 
9 
10 class _IRType(Enum):
11     """Enum to set the minimum required logging level to print a message to stdout
12     """
13     ts = 0
14     fx = 1
15 
16 
17 def _module_ir(module: Any, ir: str) -> _IRType.ts:
18     # Possible module types
19     module_is_tsable = any(
20         isinstance(module, t) for t in [torch.nn.Module, torch.jit.ScriptModule, torch.jit.ScriptFunction])
21     module_is_fxable = any(isinstance(module, t) for t in [torch.nn.Module, torch.fx.GraphModule])
22 
23     ir_targets_torchscript = any([ir == opt for opt in ["torchscript", "ts"]])
24     ir_targets_fx = ir == "fx"
25 
26     if module_is_tsable and ir_targets_torchscript:
27         return _IRType.ts
28     elif module_is_fxable and ir_targets_fx:
29         if isinstance(module, torch.fx.GraphModule):
30             raise ValueError("Was given a torch.fx.GraphModule, fx is not currently supported by Torch-TensorRT")
31         elif ir_targets_fx:
32             raise ValueError("Preferred ir was set to \"fx\" which is currently not supported by Torch-TensorRT")
33         else:
34             raise ValueError("Torch-TensorRT currently does not support fx")
35         # return _IRType.fx
36     else:
37         if ir == "default":
38             # Options are listed in order of preference
39             if module_is_tsable:
40                 logging.log(logging.Level.Info, "ir was set to default, using TorchScript as ir")
41                 return _IRType.ts
42             elif module_is_fxable:
43                 raise ValueError("Was given a torch.fx.GraphModule, fx is not currently supported by Torch-TensorRT")
44                 #logging.log(logging.Level.Info, "ir was set to default, using TorchScript as fx")
45                 #return _IRType.fx
46             else:
47                 raise ValueError("Module was provided with in an unsupported format")
48         else:
49             raise ValueError("Unknown ir was requested")
50 
51 
52 def compile(module: Any, ir="default", inputs=[], enabled_precisions=set([_enums.dtype.float]), **kwargs):
53     """Compile a PyTorch module for NVIDIA GPUs using TensorRT
54 
55     Takes a existing PyTorch module and a set of settings to configure the compiler
56     and using the path specified in ``ir`` lower and compile the module to TensorRT
57     returning a PyTorch Module back
58 
59     Converts specifically the forward method of a Module
60 
61     Arguments:
62         module (Union(torch.nn.Module,torch.jit.ScriptModule): Source module
63 
64     Keyword Arguments:
65         inputs (List[Union(torch_tensorrt.Input, torch.Tensor)]): **Required** List of specifications of input shape, dtype and memory layout for inputs to the module. This argument is required. Input Sizes can be specified as torch sizes, tuples or lists. dtypes can be specified using
66             torch datatypes or torch_tensorrt datatypes and you can use either torch devices or the torch_tensorrt device type enum
67             to select device type. ::
68 
69                 input=[
70                     torch_tensorrt.Input((1, 3, 224, 224)), # Static NCHW input shape for input #1
71                     torch_tensorrt.Input(
72                         min_shape=(1, 224, 224, 3),
73                         opt_shape=(1, 512, 512, 3),
74                         max_shape=(1, 1024, 1024, 3),
75                         dtype=torch.int32
76                         format=torch.channel_last
77                     ), # Dynamic input shape for input #2
78                     torch.randn((1, 3, 224, 244)) # Use an example tensor and let torch_tensorrt infer settings
79                 ]
80 
81         enabled_precision (Set(Union(torch.dtype, torch_tensorrt.dtype))): The set of datatypes that TensorRT can use when selecting kernels
82         ir (str): The requested strategy to compile. (Options: default - Let Torch-TensorRT decide, ts - TorchScript with scripting path)
83         **kwargs: Additional settings for the specific requested strategy (See submodules for more info)
84 
85     Returns:
86         torch.nn.Module: Compiled Module, when run it will execute via TensorRT
87     """
88     target_ir = _module_ir(module, ir)
89     if target_ir == _IRType.ts:
90         ts_mod = module
91         if isinstance(module, torch.nn.Module):
92             logging.log(
93                 logging.Level.Info,
94                 "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"
95             )
96             ts_mod = torch.jit.script(module)
97         return torch_tensorrt.ts.compile(ts_mod, inputs=inputs, enabled_precisions=enabled_precisions, **kwargs)
98     elif target_ir == _IRType.fx:
99         raise RuntimeError("fx is currently not supported")
100     else:
101         raise RuntimeError("Module is an unknown format or the ir requested is unknown")
102 
103 
104 def convert_method_to_trt_engine(module: Any,
105                                  method_name: str,
106                                  ir="default",
107                                  inputs=[],
108                                  enabled_precisions=set([_enums.dtype.float]),
109                                  **kwargs):
110     """Convert a TorchScript module method to a serialized TensorRT engine
111 
112     Converts a specified method of a module to a serialized TensorRT engine given a dictionary of conversion settings
113 
114     Arguments:
115         module (Union(torch.nn.Module,torch.jit.ScriptModule): Source module
116 
117     Keyword Arguments:
118         inputs (List[Union(torch_tensorrt.Input, torch.Tensor)]): **Required** List of specifications of input shape, dtype and memory layout for inputs to the module. This argument is required. Input Sizes can be specified as torch sizes, tuples or lists. dtypes can be specified using
119             torch datatypes or torch_tensorrt datatypes and you can use either torch devices or the torch_tensorrt device type enum
120             to select device type. ::
121 
122                 input=[
123                     torch_tensorrt.Input((1, 3, 224, 224)), # Static NCHW input shape for input #1
124                     torch_tensorrt.Input(
125                         min_shape=(1, 224, 224, 3),
126                         opt_shape=(1, 512, 512, 3),
127                         max_shape=(1, 1024, 1024, 3),
128                         dtype=torch.int32
129                         format=torch.channel_last
130                     ), # Dynamic input shape for input #2
131                     torch.randn((1, 3, 224, 244)) # Use an example tensor and let torch_tensorrt infer settings
132                 ]
133 
134         enabled_precision (Set(Union(torch.dtype, torch_tensorrt.dtype))): The set of datatypes that TensorRT can use when selecting kernels
135         ir (str): The requested strategy to compile. (Options: default - Let Torch-TensorRT decide, ts - TorchScript with scripting path)
136         **kwargs: Additional settings for the specific requested strategy (See submodules for more info)
137 
138     Returns:
139         bytes: Serialized TensorRT engine, can either be saved to a file or deserialized via TensorRT APIs
140     """
141     target_ir = _module_ir(module, ir)
142     if target_ir == _IRType.ts:
143         ts_mod = module
144         if isinstance(module, torch.nn.Module):
145             logging.log(
146                 logging.Level.Info,
147                 "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"
148             )
149             ts_mod = torch.jit.script(module)
150         return torch_tensorrt.ts.convert_method_to_trt_engine(ts_mod,
151                                                               method_name,
152                                                               inputs=inputs,
153                                                               enabled_precisions=enabled_precisions,
154                                                               **kwargs)
155     elif target_ir == _IRType.fx:
156         raise RuntimeError("fx is currently not supported")
157     else:
158         raise RuntimeError("Module is an unknown format or the ir requested is unknown")
159
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
     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 @@
     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 @@
     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 @@
         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 @@
     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

{"golden_diff": "diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py\n--- a/py/torch_tensorrt/_compile.py\n+++ b/py/torch_tensorrt/_compile.py\n@@ -14,11 +14,28 @@\n     fx = 1\n \n \n-def _module_ir(module: Any, ir: str) -> _IRType.ts:\n-    # Possible module types\n-    module_is_tsable = any(\n-        isinstance(module, t) for t in [torch.nn.Module, torch.jit.ScriptModule, torch.jit.ScriptFunction])\n-    module_is_fxable = any(isinstance(module, t) for t in [torch.nn.Module, torch.fx.GraphModule])\n+class _ModuleType(Enum):\n+    \"\"\"Enum to set the minimum required logging level to print a message to stdout\n+    \"\"\"\n+    nn = 0\n+    ts = 1\n+    fx = 2\n+\n+\n+def _parse_module_type(module: Any) -> _ModuleType:\n+    if any(isinstance(module, t) for t in [torch.jit.ScriptModule, torch.jit.ScriptFunction]):\n+        return _ModuleType.ts\n+    elif isinstance(module, torch.fx.GraphModule):\n+        return _ModuleType.fx\n+    elif isinstance(module, torch.nn.Module):\n+        return _ModuleType.nn\n+    else:\n+        raise RuntimeError(\"Module is an unknown format\")\n+\n+\n+def _get_target_ir(module_type: _ModuleType, ir: str) -> _IRType:\n+    module_is_tsable = any([module_type == t for t in [_ModuleType.nn, _ModuleType.ts]])\n+    module_is_fxable = any([module_type == t for t in [_ModuleType.nn, _ModuleType.fx]])\n \n     ir_targets_torchscript = any([ir == opt for opt in [\"torchscript\", \"ts\"]])\n     ir_targets_fx = ir == \"fx\"\n@@ -26,7 +43,7 @@\n     if module_is_tsable and ir_targets_torchscript:\n         return _IRType.ts\n     elif module_is_fxable and ir_targets_fx:\n-        if isinstance(module, torch.fx.GraphModule):\n+        if module_type == _ModuleType.fx:\n             raise ValueError(\"Was given a torch.fx.GraphModule, fx is not currently supported by Torch-TensorRT\")\n         elif ir_targets_fx:\n             raise ValueError(\"Preferred ir was set to \\\"fx\\\" which is currently not supported by Torch-TensorRT\")\n@@ -85,10 +102,11 @@\n     Returns:\n         torch.nn.Module: Compiled Module, when run it will execute via TensorRT\n     \"\"\"\n-    target_ir = _module_ir(module, ir)\n+    module_type = _parse_module_type(module)\n+    target_ir = _get_target_ir(module_type, ir)\n     if target_ir == _IRType.ts:\n         ts_mod = module\n-        if isinstance(module, torch.nn.Module):\n+        if module_type == _ModuleType.nn:\n             logging.log(\n                 logging.Level.Info,\n                 \"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\"\n@@ -134,14 +152,14 @@\n         enabled_precision (Set(Union(torch.dtype, torch_tensorrt.dtype))): The set of datatypes that TensorRT can use when selecting kernels\n         ir (str): The requested strategy to compile. (Options: default - Let Torch-TensorRT decide, ts - TorchScript with scripting path)\n         **kwargs: Additional settings for the specific requested strategy (See submodules for more info)\n-\n     Returns:\n         bytes: Serialized TensorRT engine, can either be saved to a file or deserialized via TensorRT APIs\n     \"\"\"\n-    target_ir = _module_ir(module, ir)\n+    module_type = _parse_module_type(module)\n+    target_ir = _get_target_ir(module_type, ir)\n     if target_ir == _IRType.ts:\n         ts_mod = module\n-        if isinstance(module, torch.nn.Module):\n+        if module_type == _ModuleType.nn:\n             logging.log(\n                 logging.Level.Info,\n                 \"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\"\n@@ -155,4 +173,4 @@\n     elif target_ir == _IRType.fx:\n         raise RuntimeError(\"fx is currently not supported\")\n     else:\n-        raise RuntimeError(\"Module is an unknown format or the ir requested is unknown\")\n+        raise RuntimeError(\"Module is an unknown format or the ir requested is unknown\")\n\\ No newline at end of file\n", "issue": "\ud83d\udc1b [Bug] Incorrect judgment of model type\n##  Bug Description\r\n\r\nWhen I provide a TorchScript model, Torch-TensorRT claims it is an nn.Module model\r\n\r\n```\r\nINFO: [Torch-TensorRT] - ir was set to default, using TorchScript as ir\r\nINFO: [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\r\n```\r\n\r\n<!-- A clear and concise description of what the bug is. -->\r\n\r\n## To Reproduce\r\n\r\nJust run:\r\n\r\n```\r\nimport torch\r\nimport torch.nn as nn\r\nimport torch_tensorrt as tt\r\n\r\nclass Model(nn.Module):\r\n    def __init__(self):\r\n        super(Model,self).__init__()\r\n        self.conv=nn.Conv2d(3,3,3)\r\n\r\n    def forward(self,x1):\r\n        out=self.conv(x1)\r\n        return out\r\n\r\na=Model().cuda().eval()\r\nb=torch.jit.trace(a,torch.ones([1,3,20,20]).cuda())\r\n\r\ncompile_settings = {}\r\ncompile_settings[\"inputs\"] = [tt.Input(shape = [1,3,20,20])]\r\ntt.logging.set_reportable_log_level(tt.logging.Level.Graph)\r\n\r\ntt.compile(b,**compile_settings)\r\n```\r\n\r\n<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->\r\n\r\n## Expected behavior\r\n\r\nModel type is correctly judged.\r\n\r\n<!-- A clear and concise description of what you expected to happen. -->\r\n\r\n## Environment\r\n\r\n> Build information about Torch-TensorRT can be found by turning on debug messages\r\n\r\n - Torch-TensorRT Version : 1.0.0\r\n - PyTorch Version: 1.10.0\r\n - CPU Architecture: Intel(R) Xeon(R) Platinum 8352Y CPU @ 2.20GHz\r\n - OS (e.g., Linux): CentOS 7\r\n - How you installed PyTorch (`conda`, `pip`, `libtorch`, source): pip\r\n - Build command you used (if compiling from source):\r\n - Are you using local sources or building from archives:\r\n - Python version: 3.6.8\r\n - CUDA version: 11.4\r\n - GPU models and configuration: A30\r\n - Any other relevant information:\r\n\r\n## Additional context\r\n\r\nI have fixed this bug locally. After confirming the existence of this bug, I will open a PR to fix it.\r\n\r\n<!-- Add any other context about the problem here. -->\r\n\n", "before_files": [{"content": "from typing import List, Dict, Any\nfrom torch_tensorrt import _enums\nimport torch_tensorrt.ts\nfrom torch_tensorrt import logging\nimport torch\nfrom torch import fx\nfrom enum import Enum\n\n\nclass _IRType(Enum):\n    \"\"\"Enum to set the minimum required logging level to print a message to stdout\n    \"\"\"\n    ts = 0\n    fx = 1\n\n\ndef _module_ir(module: Any, ir: str) -> _IRType.ts:\n    # Possible module types\n    module_is_tsable = any(\n        isinstance(module, t) for t in [torch.nn.Module, torch.jit.ScriptModule, torch.jit.ScriptFunction])\n    module_is_fxable = any(isinstance(module, t) for t in [torch.nn.Module, torch.fx.GraphModule])\n\n    ir_targets_torchscript = any([ir == opt for opt in [\"torchscript\", \"ts\"]])\n    ir_targets_fx = ir == \"fx\"\n\n    if module_is_tsable and ir_targets_torchscript:\n        return _IRType.ts\n    elif module_is_fxable and ir_targets_fx:\n        if isinstance(module, torch.fx.GraphModule):\n            raise ValueError(\"Was given a torch.fx.GraphModule, fx is not currently supported by Torch-TensorRT\")\n        elif ir_targets_fx:\n            raise ValueError(\"Preferred ir was set to \\\"fx\\\" which is currently not supported by Torch-TensorRT\")\n        else:\n            raise ValueError(\"Torch-TensorRT currently does not support fx\")\n        # return _IRType.fx\n    else:\n        if ir == \"default\":\n            # Options are listed in order of preference\n            if module_is_tsable:\n                logging.log(logging.Level.Info, \"ir was set to default, using TorchScript as ir\")\n                return _IRType.ts\n            elif module_is_fxable:\n                raise ValueError(\"Was given a torch.fx.GraphModule, fx is not currently supported by Torch-TensorRT\")\n                #logging.log(logging.Level.Info, \"ir was set to default, using TorchScript as fx\")\n                #return _IRType.fx\n            else:\n                raise ValueError(\"Module was provided with in an unsupported format\")\n        else:\n            raise ValueError(\"Unknown ir was requested\")\n\n\ndef compile(module: Any, ir=\"default\", inputs=[], enabled_precisions=set([_enums.dtype.float]), **kwargs):\n    \"\"\"Compile a PyTorch module for NVIDIA GPUs using TensorRT\n\n    Takes a existing PyTorch module and a set of settings to configure the compiler\n    and using the path specified in ``ir`` lower and compile the module to TensorRT\n    returning a PyTorch Module back\n\n    Converts specifically the forward method of a Module\n\n    Arguments:\n        module (Union(torch.nn.Module,torch.jit.ScriptModule): Source module\n\n    Keyword Arguments:\n        inputs (List[Union(torch_tensorrt.Input, torch.Tensor)]): **Required** List of specifications of input shape, dtype and memory layout for inputs to the module. This argument is required. Input Sizes can be specified as torch sizes, tuples or lists. dtypes can be specified using\n            torch datatypes or torch_tensorrt datatypes and you can use either torch devices or the torch_tensorrt device type enum\n            to select device type. ::\n\n                input=[\n                    torch_tensorrt.Input((1, 3, 224, 224)), # Static NCHW input shape for input #1\n                    torch_tensorrt.Input(\n                        min_shape=(1, 224, 224, 3),\n                        opt_shape=(1, 512, 512, 3),\n                        max_shape=(1, 1024, 1024, 3),\n                        dtype=torch.int32\n                        format=torch.channel_last\n                    ), # Dynamic input shape for input #2\n                    torch.randn((1, 3, 224, 244)) # Use an example tensor and let torch_tensorrt infer settings\n                ]\n\n        enabled_precision (Set(Union(torch.dtype, torch_tensorrt.dtype))): The set of datatypes that TensorRT can use when selecting kernels\n        ir (str): The requested strategy to compile. (Options: default - Let Torch-TensorRT decide, ts - TorchScript with scripting path)\n        **kwargs: Additional settings for the specific requested strategy (See submodules for more info)\n\n    Returns:\n        torch.nn.Module: Compiled Module, when run it will execute via TensorRT\n    \"\"\"\n    target_ir = _module_ir(module, ir)\n    if target_ir == _IRType.ts:\n        ts_mod = module\n        if isinstance(module, torch.nn.Module):\n            logging.log(\n                logging.Level.Info,\n                \"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\"\n            )\n            ts_mod = torch.jit.script(module)\n        return torch_tensorrt.ts.compile(ts_mod, inputs=inputs, enabled_precisions=enabled_precisions, **kwargs)\n    elif target_ir == _IRType.fx:\n        raise RuntimeError(\"fx is currently not supported\")\n    else:\n        raise RuntimeError(\"Module is an unknown format or the ir requested is unknown\")\n\n\ndef convert_method_to_trt_engine(module: Any,\n                                 method_name: str,\n                                 ir=\"default\",\n                                 inputs=[],\n                                 enabled_precisions=set([_enums.dtype.float]),\n                                 **kwargs):\n    \"\"\"Convert a TorchScript module method to a serialized TensorRT engine\n\n    Converts a specified method of a module to a serialized TensorRT engine given a dictionary of conversion settings\n\n    Arguments:\n        module (Union(torch.nn.Module,torch.jit.ScriptModule): Source module\n\n    Keyword Arguments:\n        inputs (List[Union(torch_tensorrt.Input, torch.Tensor)]): **Required** List of specifications of input shape, dtype and memory layout for inputs to the module. This argument is required. Input Sizes can be specified as torch sizes, tuples or lists. dtypes can be specified using\n            torch datatypes or torch_tensorrt datatypes and you can use either torch devices or the torch_tensorrt device type enum\n            to select device type. ::\n\n                input=[\n                    torch_tensorrt.Input((1, 3, 224, 224)), # Static NCHW input shape for input #1\n                    torch_tensorrt.Input(\n                        min_shape=(1, 224, 224, 3),\n                        opt_shape=(1, 512, 512, 3),\n                        max_shape=(1, 1024, 1024, 3),\n                        dtype=torch.int32\n                        format=torch.channel_last\n                    ), # Dynamic input shape for input #2\n                    torch.randn((1, 3, 224, 244)) # Use an example tensor and let torch_tensorrt infer settings\n                ]\n\n        enabled_precision (Set(Union(torch.dtype, torch_tensorrt.dtype))): The set of datatypes that TensorRT can use when selecting kernels\n        ir (str): The requested strategy to compile. (Options: default - Let Torch-TensorRT decide, ts - TorchScript with scripting path)\n        **kwargs: Additional settings for the specific requested strategy (See submodules for more info)\n\n    Returns:\n        bytes: Serialized TensorRT engine, can either be saved to a file or deserialized via TensorRT APIs\n    \"\"\"\n    target_ir = _module_ir(module, ir)\n    if target_ir == _IRType.ts:\n        ts_mod = module\n        if isinstance(module, torch.nn.Module):\n            logging.log(\n                logging.Level.Info,\n                \"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\"\n            )\n            ts_mod = torch.jit.script(module)\n        return torch_tensorrt.ts.convert_method_to_trt_engine(ts_mod,\n                                                              method_name,\n                                                              inputs=inputs,\n                                                              enabled_precisions=enabled_precisions,\n                                                              **kwargs)\n    elif target_ir == _IRType.fx:\n        raise RuntimeError(\"fx is currently not supported\")\n    else:\n        raise RuntimeError(\"Module is an unknown format or the ir requested is unknown\")\n", "path": "py/torch_tensorrt/_compile.py"}], "after_files": [{"content": "from typing import List, Dict, Any\nfrom torch_tensorrt import _enums\nimport torch_tensorrt.ts\nfrom torch_tensorrt import logging\nimport torch\nfrom torch import fx\nfrom enum import Enum\n\n\nclass _IRType(Enum):\n    \"\"\"Enum to set the minimum required logging level to print a message to stdout\n    \"\"\"\n    ts = 0\n    fx = 1\n\n\nclass _ModuleType(Enum):\n    \"\"\"Enum to set the minimum required logging level to print a message to stdout\n    \"\"\"\n    nn = 0\n    ts = 1\n    fx = 2\n\n\ndef _parse_module_type(module: Any) -> _ModuleType:\n    if any(isinstance(module, t) for t in [torch.jit.ScriptModule, torch.jit.ScriptFunction]):\n        return _ModuleType.ts\n    elif isinstance(module, torch.fx.GraphModule):\n        return _ModuleType.fx\n    elif isinstance(module, torch.nn.Module):\n        return _ModuleType.nn\n    else:\n        raise RuntimeError(\"Module is an unknown format\")\n\n\ndef _get_target_ir(module_type: _ModuleType, ir: str) -> _IRType:\n    module_is_tsable = any([module_type == t for t in [_ModuleType.nn, _ModuleType.ts]])\n    module_is_fxable = any([module_type == t for t in [_ModuleType.nn, _ModuleType.fx]])\n\n    ir_targets_torchscript = any([ir == opt for opt in [\"torchscript\", \"ts\"]])\n    ir_targets_fx = ir == \"fx\"\n\n    if module_is_tsable and ir_targets_torchscript:\n        return _IRType.ts\n    elif module_is_fxable and ir_targets_fx:\n        if module_type == _ModuleType.fx:\n            raise ValueError(\"Was given a torch.fx.GraphModule, fx is not currently supported by Torch-TensorRT\")\n        elif ir_targets_fx:\n            raise ValueError(\"Preferred ir was set to \\\"fx\\\" which is currently not supported by Torch-TensorRT\")\n        else:\n            raise ValueError(\"Torch-TensorRT currently does not support fx\")\n        # return _IRType.fx\n    else:\n        if ir == \"default\":\n            # Options are listed in order of preference\n            if module_is_tsable:\n                logging.log(logging.Level.Info, \"ir was set to default, using TorchScript as ir\")\n                return _IRType.ts\n            elif module_is_fxable:\n                raise ValueError(\"Was given a torch.fx.GraphModule, fx is not currently supported by Torch-TensorRT\")\n                #logging.log(logging.Level.Info, \"ir was set to default, using TorchScript as fx\")\n                #return _IRType.fx\n            else:\n                raise ValueError(\"Module was provided with in an unsupported format\")\n        else:\n            raise ValueError(\"Unknown ir was requested\")\n\n\ndef compile(module: Any, ir=\"default\", inputs=[], enabled_precisions=set([_enums.dtype.float]), **kwargs):\n    \"\"\"Compile a PyTorch module for NVIDIA GPUs using TensorRT\n\n    Takes a existing PyTorch module and a set of settings to configure the compiler\n    and using the path specified in ``ir`` lower and compile the module to TensorRT\n    returning a PyTorch Module back\n\n    Converts specifically the forward method of a Module\n\n    Arguments:\n        module (Union(torch.nn.Module,torch.jit.ScriptModule): Source module\n\n    Keyword Arguments:\n        inputs (List[Union(torch_tensorrt.Input, torch.Tensor)]): **Required** List of specifications of input shape, dtype and memory layout for inputs to the module. This argument is required. Input Sizes can be specified as torch sizes, tuples or lists. dtypes can be specified using\n            torch datatypes or torch_tensorrt datatypes and you can use either torch devices or the torch_tensorrt device type enum\n            to select device type. ::\n\n                input=[\n                    torch_tensorrt.Input((1, 3, 224, 224)), # Static NCHW input shape for input #1\n                    torch_tensorrt.Input(\n                        min_shape=(1, 224, 224, 3),\n                        opt_shape=(1, 512, 512, 3),\n                        max_shape=(1, 1024, 1024, 3),\n                        dtype=torch.int32\n                        format=torch.channel_last\n                    ), # Dynamic input shape for input #2\n                    torch.randn((1, 3, 224, 244)) # Use an example tensor and let torch_tensorrt infer settings\n                ]\n\n        enabled_precision (Set(Union(torch.dtype, torch_tensorrt.dtype))): The set of datatypes that TensorRT can use when selecting kernels\n        ir (str): The requested strategy to compile. (Options: default - Let Torch-TensorRT decide, ts - TorchScript with scripting path)\n        **kwargs: Additional settings for the specific requested strategy (See submodules for more info)\n\n    Returns:\n        torch.nn.Module: Compiled Module, when run it will execute via TensorRT\n    \"\"\"\n    module_type = _parse_module_type(module)\n    target_ir = _get_target_ir(module_type, ir)\n    if target_ir == _IRType.ts:\n        ts_mod = module\n        if module_type == _ModuleType.nn:\n            logging.log(\n                logging.Level.Info,\n                \"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\"\n            )\n            ts_mod = torch.jit.script(module)\n        return torch_tensorrt.ts.compile(ts_mod, inputs=inputs, enabled_precisions=enabled_precisions, **kwargs)\n    elif target_ir == _IRType.fx:\n        raise RuntimeError(\"fx is currently not supported\")\n    else:\n        raise RuntimeError(\"Module is an unknown format or the ir requested is unknown\")\n\n\ndef convert_method_to_trt_engine(module: Any,\n                                 method_name: str,\n                                 ir=\"default\",\n                                 inputs=[],\n                                 enabled_precisions=set([_enums.dtype.float]),\n                                 **kwargs):\n    \"\"\"Convert a TorchScript module method to a serialized TensorRT engine\n\n    Converts a specified method of a module to a serialized TensorRT engine given a dictionary of conversion settings\n\n    Arguments:\n        module (Union(torch.nn.Module,torch.jit.ScriptModule): Source module\n\n    Keyword Arguments:\n        inputs (List[Union(torch_tensorrt.Input, torch.Tensor)]): **Required** List of specifications of input shape, dtype and memory layout for inputs to the module. This argument is required. Input Sizes can be specified as torch sizes, tuples or lists. dtypes can be specified using\n            torch datatypes or torch_tensorrt datatypes and you can use either torch devices or the torch_tensorrt device type enum\n            to select device type. ::\n\n                input=[\n                    torch_tensorrt.Input((1, 3, 224, 224)), # Static NCHW input shape for input #1\n                    torch_tensorrt.Input(\n                        min_shape=(1, 224, 224, 3),\n                        opt_shape=(1, 512, 512, 3),\n                        max_shape=(1, 1024, 1024, 3),\n                        dtype=torch.int32\n                        format=torch.channel_last\n                    ), # Dynamic input shape for input #2\n                    torch.randn((1, 3, 224, 244)) # Use an example tensor and let torch_tensorrt infer settings\n                ]\n\n        enabled_precision (Set(Union(torch.dtype, torch_tensorrt.dtype))): The set of datatypes that TensorRT can use when selecting kernels\n        ir (str): The requested strategy to compile. (Options: default - Let Torch-TensorRT decide, ts - TorchScript with scripting path)\n        **kwargs: Additional settings for the specific requested strategy (See submodules for more info)\n    Returns:\n        bytes: Serialized TensorRT engine, can either be saved to a file or deserialized via TensorRT APIs\n    \"\"\"\n    module_type = _parse_module_type(module)\n    target_ir = _get_target_ir(module_type, ir)\n    if target_ir == _IRType.ts:\n        ts_mod = module\n        if module_type == _ModuleType.nn:\n            logging.log(\n                logging.Level.Info,\n                \"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\"\n            )\n            ts_mod = torch.jit.script(module)\n        return torch_tensorrt.ts.convert_method_to_trt_engine(ts_mod,\n                                                              method_name,\n                                                              inputs=inputs,\n                                                              enabled_precisions=enabled_precisions,\n                                                              **kwargs)\n    elif target_ir == _IRType.fx:\n        raise RuntimeError(\"fx is currently not supported\")\n    else:\n        raise RuntimeError(\"Module is an unknown format or the ir requested is unknown\")", "path": "py/torch_tensorrt/_compile.py"}]}