We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
🐛 [Bug] An error occurs in CompileGraph when gpu_id == 1
When I tried to Complie on the second GPU in a multi-GPU environment, an error occurred. The code sample used is as follows.

```cpp
void load(const std::string& model_path, int64_t gpu_id, int64_t opt_batch_size) {
    torch::jit::Module module = torch::jit::load(model_path);
    torch::Device device = (torch::cuda::is_available() ? torch::Device(torch::kCUDA, gpu_id) : torch::Device(torch::kCPU));
    module.to(device, torch::kHalf);
    module.eval();

    std::vector<int64_t> in_opt = { opt_batch_size, INPUT_CHANNEL_NUM, BOARD_WIDTH, BOARD_WIDTH };

    trtorch::CompileSpec::InputRange range(in_opt);
    trtorch::CompileSpec info({ range });
    info.op_precision = torch::kHalf;
    info.device.gpu_id = gpu_id;
    module = trtorch::CompileGraph(module, info);
}
```

#### Error1
I called this function with gpu_id = 1. I got the following error:

```
terminate called after throwing an instance of 'trtorch::Error'
  what():  [enforce fail at core/conversion/conversionctx/ConversionCtx.cpp:107] Expected cudaSetDevice(settings.device.gpu_id) to be true but got false
Unable to set gpu id: 1
```

I think this line is the cause.
https://github.com/NVIDIA/TRTorch/blob/1d4b967a28e36beee048703f5645ee6fcc95793d/core/conversion/conversionctx/ConversionCtx.cpp#L112

`cudaSetDevice` returns `cudaSuccess` (= 0) on success. However, `TRTORCH_CHECK` judges success or failure as a Boolean type.

I fixed it as follows and rebuilt it so that this error disappeared.

```diff
diff --git a/core/conversion/conversionctx/ConversionCtx.cpp b/core/conversion/conversionctx/ConversionCtx.cpp
index ff23692..bc5bf68 100644
--- a/core/conversion/conversionctx/ConversionCtx.cpp
+++ b/core/conversion/conversionctx/ConversionCtx.cpp
@@ -109,7 +109,7 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)
   cfg->setEngineCapability(settings.capability);
 
   if (settings.device.gpu_id) {
-    TRTORCH_CHECK(cudaSetDevice(settings.device.gpu_id), "Unable to set gpu id: " << settings.device.gpu_id);
+    TRTORCH_CHECK(cudaSetDevice(settings.device.gpu_id) == cudaSuccess, "Unable to set gpu id: " << settings.device.gpu_id);
   }
 
   if (settings.device.device_type == nvinfer1::DeviceType::kDLA) {
```

You may also use `set_device`.

https://github.com/NVIDIA/TRTorch/blob/1d4b967a28e36beee048703f5645ee6fcc95793d/core/compiler.cpp#L176-L178

#### Error2
After making the above fix, I get the following error:

```
ERROR: [TRTorch Conversion Context] - Builder was created on device different than current device.
```

I changed `cudaSetDevice` to do it at the beginning of the function and it worked fine.

```diff
diff --git a/core/conversion/conversionctx/ConversionCtx.cpp b/core/conversion/conversionctx/ConversionCtx.cpp
index ff23692..09a419c 100644
--- a/core/conversion/conversionctx/ConversionCtx.cpp
+++ b/core/conversion/conversionctx/ConversionCtx.cpp
@@ -47,6 +47,10 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)
           util::logging::get_logger().get_reportable_severity(),
           util::logging::get_logger().get_is_colored_output_on()) {
   // TODO: Support FP16 and FP32 from JIT information
+  if (settings.device.gpu_id) {
+    TRTORCH_CHECK(cudaSetDevice(settings.device.gpu_id) == cudaSuccess, "Unable to set gpu id: " << settings.device.gpu_id);
+  }
+
   builder = nvinfer1::createInferBuilder(logger);
   net = builder->createNetworkV2(1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH));
 
@@ -108,10 +112,6 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)
   cfg->setDefaultDeviceType(settings.device.device_type);
   cfg->setEngineCapability(settings.capability);
 
-  if (settings.device.gpu_id) {
-    TRTORCH_CHECK(cudaSetDevice(settings.device.gpu_id), "Unable to set gpu id: " << settings.device.gpu_id);
-  }
-
   if (settings.device.device_type == nvinfer1::DeviceType::kDLA) {
     auto nbDLACores = builder->getNbDLACores();
     TRTORCH_CHECK(
```

It's working, but I'm not sure if this is a good fix as there may be other side effects as well.
I would appreciate it if you could respond appropriately.

--- 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/trtorch/_compiler.py`
Content:
```
1 from typing import List, Dict, Any
2 import torch
3 from torch import nn
4 
5 import trtorch._C
6 from trtorch._compile_spec import _parse_compile_spec
7 from trtorch._version import __version__
8 from types import FunctionType
9 
10 
11 def compile(module: torch.jit.ScriptModule, compile_spec: Any) -> torch.jit.ScriptModule:
12     """Compile a TorchScript module for NVIDIA GPUs using TensorRT
13 
14     Takes a existing TorchScript module and a set of settings to configure the compiler
15     and will convert methods to JIT Graphs which call equivalent TensorRT engines
16 
17     Converts specifically the forward method of a TorchScript Module
18 
19     Args:
20         module (torch.jit.ScriptModule): Source module, a result of tracing or scripting a PyTorch
21             ``torch.nn.Module``
22         compile_spec (dict): Compilation settings including operating precision, target device, etc.
23             One key is required which is ``input_shapes``, describing the input sizes or ranges for inputs
24             to the graph. All other keys are optional
25 
26             .. code-block:: py
27 
28                 compile_spec = {
29                     "input_shapes": [
30                         (1, 3, 224, 224), # Static input shape for input #1
31                         {
32                             "min": (1, 3, 224, 224),
33                             "opt": (1, 3, 512, 512),
34                             "max": (1, 3, 1024, 1024)
35                         } # Dynamic input shape for input #2
36                     ],
37                     "device": {
38                         "device_type": torch.device("cuda"), # Type of device to run engine on (for DLA use trtorch.DeviceType.DLA)
39                         "gpu_id": 0, # Target gpu id to run engine (Use Xavier as gpu id for DLA)
40                         "dla_core": 0, # (DLA only) Target dla core id to run engine
41                         "allow_gpu_fallback": false, # (DLA only) Allow layers unsupported on DLA to run on GPU
42                     },
43                     "op_precision": torch.half, # Operating precision set to FP16
44                     "refit": false, # enable refit
45                     "debug": false, # enable debuggable engine
46                     "strict_types": false, # kernels should strictly run in operating precision
47                     "capability": trtorch.EngineCapability.DEFAULT, # Restrict kernel selection to safe gpu kernels or safe dla kernels
48                     "num_min_timing_iters": 2, # Number of minimization timing iterations used to select kernels
49                     "num_avg_timing_iters": 1, # Number of averaging timing iterations used to select kernels
50                     "workspace_size": 0, # Maximum size of workspace given to TensorRT
51                     "max_batch_size": 0, # Maximum batch size (must be >= 1 to be set, 0 means not set)
52                 }
53 
54             Input Sizes can be specified as torch sizes, tuples or lists. Op precisions can be specified using
55             torch datatypes or trtorch datatypes and you can use either torch devices or the trtorch device type enum
56             to select device type.
57 
58     Returns:
59         torch.jit.ScriptModule: Compiled TorchScript Module, when run it will execute via TensorRT
60     """
61 
62     if isinstance(module, torch.jit.ScriptFunction):
63         raise TypeError(
64             "torch.jit.ScriptFunction currently is not directly supported, wrap the function in a module to compile")
65 
66     compiled_cpp_mod = trtorch._C.compile_graph(module._c, _parse_compile_spec(compile_spec))
67     compiled_module = torch.jit._recursive.wrap_cpp_module(compiled_cpp_mod)
68     return compiled_module
69 
70 
71 def convert_method_to_trt_engine(module: torch.jit.ScriptModule, method_name: str, compile_spec: Any) -> str:
72     """Convert a TorchScript module method to a serialized TensorRT engine
73 
74     Converts a specified method of a module to a serialized TensorRT engine given a dictionary of conversion settings
75 
76     Args:
77         module (torch.jit.ScriptModule): Source module, a result of tracing or scripting a PyTorch
78             ``torch.nn.Module``
79         method_name (str): Name of method to convert
80         compile_spec (dict): Compilation settings including operating precision, target device, etc.
81             One key is required which is ``input_shapes``, describing the input sizes or ranges for inputs
82             to the graph. All other keys are optional
83 
84             .. code-block:: py
85 
86                 CompileSpec = {
87                     "input_shapes": [
88                         (1, 3, 224, 224), # Static input shape for input #1
89                         {
90                             "min": (1, 3, 224, 224),
91                             "opt": (1, 3, 512, 512),
92                             "max": (1, 3, 1024, 1024)
93                         } # Dynamic input shape for input #2
94                     ],
95                     "device": {
96                         "device_type": torch.device("cuda"), # Type of device to run engine on (for DLA use trtorch.DeviceType.DLA)
97                         "gpu_id": 0, # Target gpu id to run engine (Use Xavier as gpu id for DLA)
98                         "dla_core": 0, # (DLA only) Target dla core id to run engine
99                         "allow_gpu_fallback": false, # (DLA only) Allow layers unsupported on DLA to run on GPU
100                     },
101                     "op_precision": torch.half, # Operating precision set to FP16
102                     "disable_tf32": False, # Force FP32 layers to use traditional as FP32 format vs the default behavior of rounding the inputs to 10-bit mantissas before multiplying, but accumulates the sum using 23-bit mantissas
103                     "refit": false, # enable refit
104                     "debug": false, # enable debuggable engine
105                     "strict_types": false, # kernels should strictly run in operating precision
106                     "capability": trtorch.EngineCapability.DEFAULT, # Restrict kernel selection to safe gpu kernels or safe dla kernels
107                     "num_min_timing_iters": 2, # Number of minimization timing iterations used to select kernels
108                     "num_avg_timing_iters": 1, # Number of averaging timing iterations used to select kernels
109                     "workspace_size": 0, # Maximum size of workspace given to TensorRT
110                     "max_batch_size": 0, # Maximum batch size (must be >= 1 to be set, 0 means not set)
111                 }
112 
113             Input Sizes can be specified as torch sizes, tuples or lists. Op precisions can be specified using
114             torch datatypes or trtorch datatypes and you can use either torch devices or the trtorch device type enum
115             to select device type.
116 
117     Returns:
118         bytes: Serialized TensorRT engine, can either be saved to a file or deserialized via TensorRT APIs
119     """
120     if isinstance(module, torch.jit.ScriptFunction):
121         raise TypeError(
122             "torch.jit.ScriptFunctions currently are not directly supported, wrap the function in a module to compile")
123 
124     return trtorch._C.convert_graph_to_trt_engine(module._c, method_name, _parse_compile_spec(compile_spec))
125 
126 
127 def check_method_op_support(module: torch.jit.ScriptModule, method_name: str) -> bool:
128     """Checks to see if a method is fully supported by TRTorch
129 
130     Checks if a method of a TorchScript module can be compiled by TRTorch, if not, a list of operators
131     that are not supported are printed out and the function returns false, else true.
132 
133     Args:
134         module (torch.jit.ScriptModule): Source module, a result of tracing or scripting a PyTorch
135             ``torch.nn.Module``
136         method_name (str): Name of method to check
137 
138     Returns:
139         bool: True if supported Method
140     """
141     return trtorch._C.check_method_op_support(module._c, method_name)
142 
143 
144 def dump_build_info():
145     """Prints build information about the TRTorch distribution to stdout
146     """
147     print(get_build_info())
148 
149 
150 def get_build_info() -> str:
151     """Returns a string containing the build information of TRTorch distribution
152 
153     Returns:
154         str: String containing the build information for TRTorch distribution
155     """
156     build_info = trtorch._C.get_build_info()
157     build_info = "TRTorch Version: " + str(__version__) + '\n' + build_info
158     return build_info
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/trtorch/_compiler.py b/py/trtorch/_compiler.py
--- a/py/trtorch/_compiler.py
+++ b/py/trtorch/_compiler.py
@@ -156,3 +156,6 @@
     build_info = trtorch._C.get_build_info()
     build_info = "TRTorch Version: " + str(__version__) + '\n' + build_info
     return build_info
+
+def set_device(gpu_id):
+    trtorch._C.set_device(gpu_id)

{"golden_diff": "diff --git a/py/trtorch/_compiler.py b/py/trtorch/_compiler.py\n--- a/py/trtorch/_compiler.py\n+++ b/py/trtorch/_compiler.py\n@@ -156,3 +156,6 @@\n     build_info = trtorch._C.get_build_info()\n     build_info = \"TRTorch Version: \" + str(__version__) + '\\n' + build_info\n     return build_info\n+\n+def set_device(gpu_id):\n+    trtorch._C.set_device(gpu_id)\n", "issue": "\ud83d\udc1b [Bug] An error occurs in CompileGraph when gpu_id == 1\nWhen I tried to Complie on the second GPU in a multi-GPU environment, an error occurred. The code sample used is as follows.\r\n\r\n```cpp\r\nvoid load(const std::string& model_path, int64_t gpu_id, int64_t opt_batch_size) {\r\n    torch::jit::Module module = torch::jit::load(model_path);\r\n    torch::Device device = (torch::cuda::is_available() ? torch::Device(torch::kCUDA, gpu_id) : torch::Device(torch::kCPU));\r\n    module.to(device, torch::kHalf);\r\n    module.eval();\r\n\r\n    std::vector<int64_t> in_opt = { opt_batch_size, INPUT_CHANNEL_NUM, BOARD_WIDTH, BOARD_WIDTH };\r\n\r\n    trtorch::CompileSpec::InputRange range(in_opt);\r\n    trtorch::CompileSpec info({ range });\r\n    info.op_precision = torch::kHalf;\r\n    info.device.gpu_id = gpu_id;\r\n    module = trtorch::CompileGraph(module, info);\r\n}\r\n```\r\n\r\n#### Error1\r\nI called this function with gpu_id = 1. I got the following error:\r\n\r\n```\r\nterminate called after throwing an instance of 'trtorch::Error'\r\n  what():  [enforce fail at core/conversion/conversionctx/ConversionCtx.cpp:107] Expected cudaSetDevice(settings.device.gpu_id) to be true but got false\r\nUnable to set gpu id: 1\r\n```\r\n\r\nI think this line is the cause.\r\nhttps://github.com/NVIDIA/TRTorch/blob/1d4b967a28e36beee048703f5645ee6fcc95793d/core/conversion/conversionctx/ConversionCtx.cpp#L112\r\n\r\n`cudaSetDevice` returns `cudaSuccess` (= 0) on success. However, `TRTORCH_CHECK` judges success or failure as a Boolean type.\r\n\r\nI fixed it as follows and rebuilt it so that this error disappeared.\r\n\r\n```diff\r\ndiff --git a/core/conversion/conversionctx/ConversionCtx.cpp b/core/conversion/conversionctx/ConversionCtx.cpp\r\nindex ff23692..bc5bf68 100644\r\n--- a/core/conversion/conversionctx/ConversionCtx.cpp\r\n+++ b/core/conversion/conversionctx/ConversionCtx.cpp\r\n@@ -109,7 +109,7 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)\r\n   cfg->setEngineCapability(settings.capability);\r\n \r\n   if (settings.device.gpu_id) {\r\n-    TRTORCH_CHECK(cudaSetDevice(settings.device.gpu_id), \"Unable to set gpu id: \" << settings.device.gpu_id);\r\n+    TRTORCH_CHECK(cudaSetDevice(settings.device.gpu_id) == cudaSuccess, \"Unable to set gpu id: \" << settings.device.gpu_id);\r\n   }\r\n \r\n   if (settings.device.device_type == nvinfer1::DeviceType::kDLA) {\r\n```\r\n\r\nYou may also use `set_device`.\r\n\r\nhttps://github.com/NVIDIA/TRTorch/blob/1d4b967a28e36beee048703f5645ee6fcc95793d/core/compiler.cpp#L176-L178\r\n\r\n#### Error2\r\nAfter making the above fix, I get the following error:\r\n\r\n```\r\nERROR: [TRTorch Conversion Context] - Builder was created on device different than current device.\r\n```\r\n\r\nI changed `cudaSetDevice` to do it at the beginning of the function and it worked fine.\r\n\r\n```diff\r\ndiff --git a/core/conversion/conversionctx/ConversionCtx.cpp b/core/conversion/conversionctx/ConversionCtx.cpp\r\nindex ff23692..09a419c 100644\r\n--- a/core/conversion/conversionctx/ConversionCtx.cpp\r\n+++ b/core/conversion/conversionctx/ConversionCtx.cpp\r\n@@ -47,6 +47,10 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)\r\n           util::logging::get_logger().get_reportable_severity(),\r\n           util::logging::get_logger().get_is_colored_output_on()) {\r\n   // TODO: Support FP16 and FP32 from JIT information\r\n+  if (settings.device.gpu_id) {\r\n+    TRTORCH_CHECK(cudaSetDevice(settings.device.gpu_id) == cudaSuccess, \"Unable to set gpu id: \" << settings.device.gpu_id);\r\n+  }\r\n+\r\n   builder = nvinfer1::createInferBuilder(logger);\r\n   net = builder->createNetworkV2(1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH));\r\n \r\n@@ -108,10 +112,6 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)\r\n   cfg->setDefaultDeviceType(settings.device.device_type);\r\n   cfg->setEngineCapability(settings.capability);\r\n \r\n-  if (settings.device.gpu_id) {\r\n-    TRTORCH_CHECK(cudaSetDevice(settings.device.gpu_id), \"Unable to set gpu id: \" << settings.device.gpu_id);\r\n-  }\r\n-\r\n   if (settings.device.device_type == nvinfer1::DeviceType::kDLA) {\r\n     auto nbDLACores = builder->getNbDLACores();\r\n     TRTORCH_CHECK(\r\n```\r\n\r\nIt's working, but I'm not sure if this is a good fix as there may be other side effects as well.\r\nI would appreciate it if you could respond appropriately.\n", "before_files": [{"content": "from typing import List, Dict, Any\nimport torch\nfrom torch import nn\n\nimport trtorch._C\nfrom trtorch._compile_spec import _parse_compile_spec\nfrom trtorch._version import __version__\nfrom types import FunctionType\n\n\ndef compile(module: torch.jit.ScriptModule, compile_spec: Any) -> torch.jit.ScriptModule:\n    \"\"\"Compile a TorchScript module for NVIDIA GPUs using TensorRT\n\n    Takes a existing TorchScript module and a set of settings to configure the compiler\n    and will convert methods to JIT Graphs which call equivalent TensorRT engines\n\n    Converts specifically the forward method of a TorchScript Module\n\n    Args:\n        module (torch.jit.ScriptModule): Source module, a result of tracing or scripting a PyTorch\n            ``torch.nn.Module``\n        compile_spec (dict): Compilation settings including operating precision, target device, etc.\n            One key is required which is ``input_shapes``, describing the input sizes or ranges for inputs\n            to the graph. All other keys are optional\n\n            .. code-block:: py\n\n                compile_spec = {\n                    \"input_shapes\": [\n                        (1, 3, 224, 224), # Static input shape for input #1\n                        {\n                            \"min\": (1, 3, 224, 224),\n                            \"opt\": (1, 3, 512, 512),\n                            \"max\": (1, 3, 1024, 1024)\n                        } # Dynamic input shape for input #2\n                    ],\n                    \"device\": {\n                        \"device_type\": torch.device(\"cuda\"), # Type of device to run engine on (for DLA use trtorch.DeviceType.DLA)\n                        \"gpu_id\": 0, # Target gpu id to run engine (Use Xavier as gpu id for DLA)\n                        \"dla_core\": 0, # (DLA only) Target dla core id to run engine\n                        \"allow_gpu_fallback\": false, # (DLA only) Allow layers unsupported on DLA to run on GPU\n                    },\n                    \"op_precision\": torch.half, # Operating precision set to FP16\n                    \"refit\": false, # enable refit\n                    \"debug\": false, # enable debuggable engine\n                    \"strict_types\": false, # kernels should strictly run in operating precision\n                    \"capability\": trtorch.EngineCapability.DEFAULT, # Restrict kernel selection to safe gpu kernels or safe dla kernels\n                    \"num_min_timing_iters\": 2, # Number of minimization timing iterations used to select kernels\n                    \"num_avg_timing_iters\": 1, # Number of averaging timing iterations used to select kernels\n                    \"workspace_size\": 0, # Maximum size of workspace given to TensorRT\n                    \"max_batch_size\": 0, # Maximum batch size (must be >= 1 to be set, 0 means not set)\n                }\n\n            Input Sizes can be specified as torch sizes, tuples or lists. Op precisions can be specified using\n            torch datatypes or trtorch datatypes and you can use either torch devices or the trtorch device type enum\n            to select device type.\n\n    Returns:\n        torch.jit.ScriptModule: Compiled TorchScript Module, when run it will execute via TensorRT\n    \"\"\"\n\n    if isinstance(module, torch.jit.ScriptFunction):\n        raise TypeError(\n            \"torch.jit.ScriptFunction currently is not directly supported, wrap the function in a module to compile\")\n\n    compiled_cpp_mod = trtorch._C.compile_graph(module._c, _parse_compile_spec(compile_spec))\n    compiled_module = torch.jit._recursive.wrap_cpp_module(compiled_cpp_mod)\n    return compiled_module\n\n\ndef convert_method_to_trt_engine(module: torch.jit.ScriptModule, method_name: str, compile_spec: Any) -> str:\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    Args:\n        module (torch.jit.ScriptModule): Source module, a result of tracing or scripting a PyTorch\n            ``torch.nn.Module``\n        method_name (str): Name of method to convert\n        compile_spec (dict): Compilation settings including operating precision, target device, etc.\n            One key is required which is ``input_shapes``, describing the input sizes or ranges for inputs\n            to the graph. All other keys are optional\n\n            .. code-block:: py\n\n                CompileSpec = {\n                    \"input_shapes\": [\n                        (1, 3, 224, 224), # Static input shape for input #1\n                        {\n                            \"min\": (1, 3, 224, 224),\n                            \"opt\": (1, 3, 512, 512),\n                            \"max\": (1, 3, 1024, 1024)\n                        } # Dynamic input shape for input #2\n                    ],\n                    \"device\": {\n                        \"device_type\": torch.device(\"cuda\"), # Type of device to run engine on (for DLA use trtorch.DeviceType.DLA)\n                        \"gpu_id\": 0, # Target gpu id to run engine (Use Xavier as gpu id for DLA)\n                        \"dla_core\": 0, # (DLA only) Target dla core id to run engine\n                        \"allow_gpu_fallback\": false, # (DLA only) Allow layers unsupported on DLA to run on GPU\n                    },\n                    \"op_precision\": torch.half, # Operating precision set to FP16\n                    \"disable_tf32\": False, # Force FP32 layers to use traditional as FP32 format vs the default behavior of rounding the inputs to 10-bit mantissas before multiplying, but accumulates the sum using 23-bit mantissas\n                    \"refit\": false, # enable refit\n                    \"debug\": false, # enable debuggable engine\n                    \"strict_types\": false, # kernels should strictly run in operating precision\n                    \"capability\": trtorch.EngineCapability.DEFAULT, # Restrict kernel selection to safe gpu kernels or safe dla kernels\n                    \"num_min_timing_iters\": 2, # Number of minimization timing iterations used to select kernels\n                    \"num_avg_timing_iters\": 1, # Number of averaging timing iterations used to select kernels\n                    \"workspace_size\": 0, # Maximum size of workspace given to TensorRT\n                    \"max_batch_size\": 0, # Maximum batch size (must be >= 1 to be set, 0 means not set)\n                }\n\n            Input Sizes can be specified as torch sizes, tuples or lists. Op precisions can be specified using\n            torch datatypes or trtorch datatypes and you can use either torch devices or the trtorch device type enum\n            to select device type.\n\n    Returns:\n        bytes: Serialized TensorRT engine, can either be saved to a file or deserialized via TensorRT APIs\n    \"\"\"\n    if isinstance(module, torch.jit.ScriptFunction):\n        raise TypeError(\n            \"torch.jit.ScriptFunctions currently are not directly supported, wrap the function in a module to compile\")\n\n    return trtorch._C.convert_graph_to_trt_engine(module._c, method_name, _parse_compile_spec(compile_spec))\n\n\ndef check_method_op_support(module: torch.jit.ScriptModule, method_name: str) -> bool:\n    \"\"\"Checks to see if a method is fully supported by TRTorch\n\n    Checks if a method of a TorchScript module can be compiled by TRTorch, if not, a list of operators\n    that are not supported are printed out and the function returns false, else true.\n\n    Args:\n        module (torch.jit.ScriptModule): Source module, a result of tracing or scripting a PyTorch\n            ``torch.nn.Module``\n        method_name (str): Name of method to check\n\n    Returns:\n        bool: True if supported Method\n    \"\"\"\n    return trtorch._C.check_method_op_support(module._c, method_name)\n\n\ndef dump_build_info():\n    \"\"\"Prints build information about the TRTorch distribution to stdout\n    \"\"\"\n    print(get_build_info())\n\n\ndef get_build_info() -> str:\n    \"\"\"Returns a string containing the build information of TRTorch distribution\n\n    Returns:\n        str: String containing the build information for TRTorch distribution\n    \"\"\"\n    build_info = trtorch._C.get_build_info()\n    build_info = \"TRTorch Version: \" + str(__version__) + '\\n' + build_info\n    return build_info\n", "path": "py/trtorch/_compiler.py"}], "after_files": [{"content": "from typing import List, Dict, Any\nimport torch\nfrom torch import nn\n\nimport trtorch._C\nfrom trtorch._compile_spec import _parse_compile_spec\nfrom trtorch._version import __version__\nfrom types import FunctionType\n\n\ndef compile(module: torch.jit.ScriptModule, compile_spec: Any) -> torch.jit.ScriptModule:\n    \"\"\"Compile a TorchScript module for NVIDIA GPUs using TensorRT\n\n    Takes a existing TorchScript module and a set of settings to configure the compiler\n    and will convert methods to JIT Graphs which call equivalent TensorRT engines\n\n    Converts specifically the forward method of a TorchScript Module\n\n    Args:\n        module (torch.jit.ScriptModule): Source module, a result of tracing or scripting a PyTorch\n            ``torch.nn.Module``\n        compile_spec (dict): Compilation settings including operating precision, target device, etc.\n            One key is required which is ``input_shapes``, describing the input sizes or ranges for inputs\n            to the graph. All other keys are optional\n\n            .. code-block:: py\n\n                compile_spec = {\n                    \"input_shapes\": [\n                        (1, 3, 224, 224), # Static input shape for input #1\n                        {\n                            \"min\": (1, 3, 224, 224),\n                            \"opt\": (1, 3, 512, 512),\n                            \"max\": (1, 3, 1024, 1024)\n                        } # Dynamic input shape for input #2\n                    ],\n                    \"device\": {\n                        \"device_type\": torch.device(\"cuda\"), # Type of device to run engine on (for DLA use trtorch.DeviceType.DLA)\n                        \"gpu_id\": 0, # Target gpu id to run engine (Use Xavier as gpu id for DLA)\n                        \"dla_core\": 0, # (DLA only) Target dla core id to run engine\n                        \"allow_gpu_fallback\": false, # (DLA only) Allow layers unsupported on DLA to run on GPU\n                    },\n                    \"op_precision\": torch.half, # Operating precision set to FP16\n                    \"refit\": false, # enable refit\n                    \"debug\": false, # enable debuggable engine\n                    \"strict_types\": false, # kernels should strictly run in operating precision\n                    \"capability\": trtorch.EngineCapability.DEFAULT, # Restrict kernel selection to safe gpu kernels or safe dla kernels\n                    \"num_min_timing_iters\": 2, # Number of minimization timing iterations used to select kernels\n                    \"num_avg_timing_iters\": 1, # Number of averaging timing iterations used to select kernels\n                    \"workspace_size\": 0, # Maximum size of workspace given to TensorRT\n                    \"max_batch_size\": 0, # Maximum batch size (must be >= 1 to be set, 0 means not set)\n                }\n\n            Input Sizes can be specified as torch sizes, tuples or lists. Op precisions can be specified using\n            torch datatypes or trtorch datatypes and you can use either torch devices or the trtorch device type enum\n            to select device type.\n\n    Returns:\n        torch.jit.ScriptModule: Compiled TorchScript Module, when run it will execute via TensorRT\n    \"\"\"\n\n    if isinstance(module, torch.jit.ScriptFunction):\n        raise TypeError(\n            \"torch.jit.ScriptFunction currently is not directly supported, wrap the function in a module to compile\")\n\n    compiled_cpp_mod = trtorch._C.compile_graph(module._c, _parse_compile_spec(compile_spec))\n    compiled_module = torch.jit._recursive.wrap_cpp_module(compiled_cpp_mod)\n    return compiled_module\n\n\ndef convert_method_to_trt_engine(module: torch.jit.ScriptModule, method_name: str, compile_spec: Any) -> str:\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    Args:\n        module (torch.jit.ScriptModule): Source module, a result of tracing or scripting a PyTorch\n            ``torch.nn.Module``\n        method_name (str): Name of method to convert\n        compile_spec (dict): Compilation settings including operating precision, target device, etc.\n            One key is required which is ``input_shapes``, describing the input sizes or ranges for inputs\n            to the graph. All other keys are optional\n\n            .. code-block:: py\n\n                CompileSpec = {\n                    \"input_shapes\": [\n                        (1, 3, 224, 224), # Static input shape for input #1\n                        {\n                            \"min\": (1, 3, 224, 224),\n                            \"opt\": (1, 3, 512, 512),\n                            \"max\": (1, 3, 1024, 1024)\n                        } # Dynamic input shape for input #2\n                    ],\n                    \"device\": {\n                        \"device_type\": torch.device(\"cuda\"), # Type of device to run engine on (for DLA use trtorch.DeviceType.DLA)\n                        \"gpu_id\": 0, # Target gpu id to run engine (Use Xavier as gpu id for DLA)\n                        \"dla_core\": 0, # (DLA only) Target dla core id to run engine\n                        \"allow_gpu_fallback\": false, # (DLA only) Allow layers unsupported on DLA to run on GPU\n                    },\n                    \"op_precision\": torch.half, # Operating precision set to FP16\n                    \"disable_tf32\": False, # Force FP32 layers to use traditional as FP32 format vs the default behavior of rounding the inputs to 10-bit mantissas before multiplying, but accumulates the sum using 23-bit mantissas\n                    \"refit\": false, # enable refit\n                    \"debug\": false, # enable debuggable engine\n                    \"strict_types\": false, # kernels should strictly run in operating precision\n                    \"capability\": trtorch.EngineCapability.DEFAULT, # Restrict kernel selection to safe gpu kernels or safe dla kernels\n                    \"num_min_timing_iters\": 2, # Number of minimization timing iterations used to select kernels\n                    \"num_avg_timing_iters\": 1, # Number of averaging timing iterations used to select kernels\n                    \"workspace_size\": 0, # Maximum size of workspace given to TensorRT\n                    \"max_batch_size\": 0, # Maximum batch size (must be >= 1 to be set, 0 means not set)\n                }\n\n            Input Sizes can be specified as torch sizes, tuples or lists. Op precisions can be specified using\n            torch datatypes or trtorch datatypes and you can use either torch devices or the trtorch device type enum\n            to select device type.\n\n    Returns:\n        bytes: Serialized TensorRT engine, can either be saved to a file or deserialized via TensorRT APIs\n    \"\"\"\n    if isinstance(module, torch.jit.ScriptFunction):\n        raise TypeError(\n            \"torch.jit.ScriptFunctions currently are not directly supported, wrap the function in a module to compile\")\n\n    return trtorch._C.convert_graph_to_trt_engine(module._c, method_name, _parse_compile_spec(compile_spec))\n\n\ndef check_method_op_support(module: torch.jit.ScriptModule, method_name: str) -> bool:\n    \"\"\"Checks to see if a method is fully supported by TRTorch\n\n    Checks if a method of a TorchScript module can be compiled by TRTorch, if not, a list of operators\n    that are not supported are printed out and the function returns false, else true.\n\n    Args:\n        module (torch.jit.ScriptModule): Source module, a result of tracing or scripting a PyTorch\n            ``torch.nn.Module``\n        method_name (str): Name of method to check\n\n    Returns:\n        bool: True if supported Method\n    \"\"\"\n    return trtorch._C.check_method_op_support(module._c, method_name)\n\n\ndef dump_build_info():\n    \"\"\"Prints build information about the TRTorch distribution to stdout\n    \"\"\"\n    print(get_build_info())\n\n\ndef get_build_info() -> str:\n    \"\"\"Returns a string containing the build information of TRTorch distribution\n\n    Returns:\n        str: String containing the build information for TRTorch distribution\n    \"\"\"\n    build_info = trtorch._C.get_build_info()\n    build_info = \"TRTorch Version: \" + str(__version__) + '\\n' + build_info\n    return build_info\n\ndef set_device(gpu_id):\n    trtorch._C.set_device(gpu_id)\n", "path": "py/trtorch/_compiler.py"}]}