diff --git a/rankers/train/training_arguments.py b/rankers/train/training_arguments.py
index 40480ba..06f2fae 100644
--- a/rankers/train/training_arguments.py
+++ b/rankers/train/training_arguments.py
@@ -2,7 +2,7 @@
 from transformers.trainer_pt_utils import AcceleratorConfig
 from transformers.utils import is_accelerate_available
 from dataclasses import field, fields, dataclass
-from typing import List, Optional, Union
+from typing import List, Optional
 from enum import Enum
 from .. import is_ir_measures_available, seed_everything