From 343e2311b48b900c07072849570f82a5b90d2aa7 Mon Sep 17 00:00:00 2001 From: Jordan Gong Date: Sun, 27 Dec 2020 15:38:18 +0800 Subject: Prepare for FVG dataset --- utils/sampler.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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) -- cgit v1.2.3