fix Tensor is not defined (#4129)

hpcaitech · Jul 3, 2023 · 8abc877 · 8abc877
1 parent 7e46bc8
commit 8abc877
Showing 1 changed file with 1 addition and 0 deletions.
diff --git 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):