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 @@ Features (go/100_enablements):
 """
 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