diff --git a/benchmark/run.py b/benchmark/run.py
index b048424..2016701 100644
--- a/benchmark/run.py
+++ b/benchmark/run.py
@@ -1,4 +1,5 @@
 import os
+import torch
 import onnxruntime
 from transformers.onnx.features import FeaturesManager
 from transformers.onnx import validate_model_outputs