diff --git a/benchmark/run.py b/benchmark/run.py
index d702f2e..69ed532 100644
--- a/benchmark/run.py
+++ b/benchmark/run.py
@@ -1,5 +1,6 @@
 import os
 import torch
+from torch import nn
 import numpy
 import onnxruntime