diff --git a/detectree2/models/train.py b/detectree2/models/train.py
index 2c692f16..eb5dbcad 100644
--- a/detectree2/models/train.py
+++ b/detectree2/models/train.py
@@ -12,7 +12,7 @@
 import re
 import time
 from pathlib import Path
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 import cv2
 import detectron2.data.transforms as T  # noqa:N812