diff --git a/core/custom_training_loop.py b/core/custom_training_loop.py
index 0241145..a8f13fb 100644
--- a/core/custom_training_loop.py
+++ b/core/custom_training_loop.py
@@ -10,7 +10,7 @@
 """
 import datetime
 import os
-from typing import Callable, Dict, Iterable, List, Mapping, Optional
+from typing import Callable, Dict, Iterable, List, Optional
 
 
 from tml.common import log_weights