diff --git a/colossalai/amp/naive_amp/_utils.py b/colossalai/amp/naive_amp/_utils.py index ad2a2ceed..7633705e1 100644 --- a/colossalai/amp/naive_amp/_utils.py +++ b/colossalai/amp/naive_amp/_utils.py @@ -1,4 +1,5 @@ from typing import List + from torch import Tensor