Addressing review comments

Author: apbose

py/torch_tensorrt/dynamo/conversion/converter_utils.py

@@ -25,36 +25,6 @@
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
 
-# nearest, linear, cubic
-class GridSamplerInterpolation:
-    def __init__(self):
-        self.interpolator_mode = None
-
-    def __call__(self, interpolator_int):
-        if interpolator_int == 0:
-            self.interpolator_mode = trt.InterpolationMode.NEAREST
-        elif interpolator_int == 1:
-            self.interpolator_mode = trt.InterpolationMode.LINEAR
-        elif interpolator_int == 2:
-            self.interpolator_mode = trt.InterpolationMode.CUBIC
-        return self.interpolator_mode
-
-
-# zeros, border, reflection
-class GridSamplerSampling:
-    def __init__(self):
-        self.sample_mode = None
-
-    def __call__(self, sample_int):
-        if sample_int == 0:
-            self.sample_mode = trt.SampleMode.FILL
-        elif sample_int == 1:
-            self.sample_mode = trt.SampleMode.CLAMP
-        elif sample_int == 2:
-            self.sample_mode = trt.SampleMode.REFLECT
-        return self.sample_mode
-
-
 def get_node_name(node: torch.fx.Node) -> str:
     # nn_module_stack preserves the call stack of pytorch nn.modules
     # The call stack contains a detailed name of the module

py/torch_tensorrt/dynamo/conversion/impl/grid.py

@@ -5,14 +5,24 @@
 from torch.fx.node import Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
 from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
-from torch_tensorrt.dynamo.conversion.converter_utils import (
-    GridSamplerInterpolation,
-    GridSamplerSampling,
-    cast_trt_tensor,
-)
+from torch_tensorrt.dynamo.conversion.converter_utils import cast_trt_tensor
 from torch_tensorrt.fx.converters.converter_utils import set_layer_name
 from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
 
+# nearest, linear, cubic
+GridSamplerInterpolationMode = {
+    0: trt.InterpolationMode.NEAREST,
+    1: trt.InterpolationMode.LINEAR,
+    2: trt.InterpolationMode.CUBIC,
+}
+
+# zeros, border, reflection
+GridSamplerSampling = {
+    0: trt.SampleMode.FILL,
+    1: trt.SampleMode.CLAMP,
+    2: trt.SampleMode.REFLECT,
+}
+
 
 def grid(
     ctx: ConversionContext,
@@ -27,18 +37,19 @@ def grid(
     output_mask: Optional[Sequence[bool]] = None,
 ) -> TRTTensor:
     grid_layer = ctx.net.add_grid_sample(input, grid)
-    interpolation_mode_trt = GridSamplerInterpolation()
-    grid_layer.interpolation_mode = interpolation_mode_trt(interpolation_mode)
-    sample_mode_trt = GridSamplerSampling()
-    grid_layer.sample_mode = sample_mode_trt(padding_mode)
+    assert interpolation_mode in GridSamplerInterpolationMode
+    grid_layer.interpolation_mode = GridSamplerInterpolationMode.get(
+        interpolation_mode, None
+    )
+    assert padding_mode in GridSamplerSampling
+    grid_layer.sample_mode = GridSamplerSampling.get(padding_mode, None)
     grid_layer.align_corners = align_corners
     set_layer_name(grid_layer, target, name + "_grid_layer", source_ir)
     if output_mask is None:
         return grid_layer.get_output(0)
+    elif output_mask[0] and output_mask[1]:
+        return (grid_layer.get_output(0), None)
+    elif output_mask[0]:
+        return grid_layer.get_output(0)
     else:
-        if output_mask[0] and output_mask[1]:
-            return (grid_layer.get_output(0), None)
-        elif output_mask[0]:
-            return grid_layer.get_output(0)
-        else:
-            return None
+        return None