diff --git a/.flake8 b/.flake8
index 19c3a9bd6..c56cd6fba 100644
--- a/.flake8
+++ b/.flake8
@@ -6,6 +6,8 @@ per-file-ignores =
     # line too long
     egs/librispeech/ASR/*/conformer.py: E501,
     egs/aishell/ASR/*/conformer.py: E501,
+    # invalid escape sequence (cause by tex formular), W605
+    icefall/utils.py: E501, W605
 
 exclude =
   .git,
diff --git a/icefall/utils.py b/icefall/utils.py
index fa8f4d334..7e504108e 100644
--- a/icefall/utils.py
+++ b/icefall/utils.py
@@ -25,13 +25,15 @@ from collections import defaultdict
 from contextlib import contextmanager
 from datetime import datetime
 from pathlib import Path
-from typing import Dict, Iterable, List, TextIO, Tuple, Union
+from typing import Dict, Iterable, List, TextIO, Optional, Tuple, Union
 
 import k2
 import k2.version
 import kaldialign
 import torch
+import torch.nn as nn
 import torch.distributed as dist
+from torch.cuda.amp import GradScaler
 from torch.utils.tensorboard import SummaryWriter
 
 Pathlike = Union[str, Path]