diff --git a/colossalai/kernel/jit/bias_dropout_add.py b/colossalai/kernel/jit/bias_dropout_add.py
index 3687dde79a08..32965c1ebd69 100644
--- a/colossalai/kernel/jit/bias_dropout_add.py
+++ b/colossalai/kernel/jit/bias_dropout_add.py
@@ -1,4 +1,5 @@
 import torch
+from torch import Tensor
 
 
 def bias_dropout_add(x, bias, residual, prob, training):