1 files changed, 2 insertions, 1 deletions
diff --git a/utils/sampler.py b/utils/sampler.py
index 0a177d1..cdf1984 100644
--- a/utils/sampler.py
+++ b/utils/sampler.py
@@ -1,5 +1,6 @@
 import random
 from collections.abc import Iterator
+from typing import Union
 
 import numpy as np
 from torch.utils import data
@@ -10,7 +11,7 @@ from utils.dataset import CASIAB
 class TripletSampler(data.Sampler):
     def __init__(
             self,
-            data_source: CASIAB,
+            data_source: Union[CASIAB],
             batch_size: tuple[int, int]
     ):
         super().__init__(data_source)