diff --git a/setup.py b/setup.py
index 88ecb46c8..d2f109802 100644
--- a/setup.py
+++ b/setup.py
@@ -52,6 +52,8 @@ setuptools.setup(
             "pytest>=6.1.1",
             "coverage>=5.3",
             "pre-commit",
+            "torch",
+            "torchvision",
             "catboost>=0.26",
             "rgf-python",
             "optuna==2.8.0",
diff --git a/test/tune/test_lexiflow.py b/test/tune/test_lexiflow.py
index 0e128547e..2d4495456 100644
--- a/test/tune/test_lexiflow.py
+++ b/test/tune/test_lexiflow.py
@@ -5,7 +5,6 @@ from flaml import tune
 import torch.nn.functional as F
 import torchvision
 import numpy as np
-import time
 from ray import tune as raytune
 
 DEVICE = torch.device("cpu")