diff --git a/tests/py/dynamo/conversion/test_argmin_aten.py b/tests/py/dynamo/conversion/test_argmin_aten.py
index c0290f943e..7065ca748c 100644
--- a/tests/py/dynamo/conversion/test_argmin_aten.py
+++ b/tests/py/dynamo/conversion/test_argmin_aten.py
@@ -2,6 +2,7 @@
 import torch.nn as nn
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt import Input
 
 from .harness import DispatchTestCase