use enforce_tensor_types decorator to cast type

Author: zewenli98

py/torch_tensorrt/dynamo/conversion/impl/shuffle.py

@@ -1,10 +1,8 @@
 from typing import Optional, Sequence, Union
 
-import numpy as np
-import torch
 from torch.fx.node import Target
 from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
-from torch_tensorrt.dynamo.conversion.converter_utils import SourceIR, get_trt_tensor
+from torch_tensorrt.dynamo.conversion.converter_utils import SourceIR
 from torch_tensorrt.fx.converters.converter_utils import set_layer_name
 from torch_tensorrt.fx.types import TRTTensor
 
@@ -14,11 +12,9 @@ def reshape(
     target: Union[Target, str],
     source_ir: Optional[SourceIR],
     name: str,
-    input: Sequence[Union[TRTTensor, torch.Tensor, np.ndarray]],
+    input: TRTTensor,
     shape: Sequence[int],
 ) -> TRTTensor:
-    if not isinstance(input, TRTTensor):
-        input = get_trt_tensor(ctx, input, f"{name}_input")
     layer = ctx.net.add_shuffle(input)
     layer.reshape_dims = tuple(shape)
     set_layer_name(layer, target, name, source_ir)