added passing of torch2trt_kwargs to conversion context (#482)

* added passing of torch2trt_kwargs to conversion context

* added passing of torch2trt_kwargs to conversion context

Author: jaybdub

torch2trt/torch2trt.py

@@ -357,12 +357,14 @@ def wrapper(*args, **kwargs):
 
 
 class ConversionContext(object):
-    def __init__(self, network, converters=CONVERTERS):
+    
+    def __init__(self, network, converters=CONVERTERS, torch2trt_kwargs=None):
         self.network = LayerNamingNetworkWrapper(self, network)
         self.lock = False
         self.method_args = None
         self.method_kwargs = None
         self.method_return = None
+        self.torch2trt_kwargs = torch2trt_kwargs
         self.hooks = [
             ConversionHook(self, key, converter)
             for key, converter in converters.items()
@@ -487,7 +489,12 @@ def torch2trt(module,
               int8_calib_dataset=None,
               int8_calib_algorithm=DEFAULT_CALIBRATION_ALGORITHM,
               int8_calib_batch_size=1,
-              use_onnx=False):
+              use_onnx=False,
+              **kwargs):
+    
+    # capture arguments to provide to context
+    kwargs.update(locals())
+    kwargs.pop('kwargs')
 
     inputs_in = inputs
 
@@ -524,7 +531,7 @@ def torch2trt(module,
 
     else:
         network = builder.create_network()
-        with ConversionContext(network) as ctx:
+        with ConversionContext(network, torch2trt_kwargs=kwargs) as ctx:
 
             ctx.add_inputs(inputs, input_names)