diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py
index 56b147b35..6cfde69bc 100644
--- a/nerfstudio/engine/trainer.py
+++ b/nerfstudio/engine/trainer.py
@@ -24,7 +24,7 @@ import time
 from dataclasses import dataclass, field
 from pathlib import Path
 from threading import Lock
-from typing import Dict, List, Literal, Optional, Tuple, Type, cast,Dict,DefaultDict
+from typing import Dict, List, Literal, Optional, Tuple, Type, cast,Union,Dict,DefaultDict
 from collections import defaultdict
 import torch
 from nerfstudio.configs.experiment_config import ExperimentConfig