From b33281425a7adc1a971677f38c140d289978b69b Mon Sep 17 00:00:00 2001 From: henryruhs Date: Fri, 21 Feb 2025 11:03:31 +0100 Subject: [PATCH] This is not a TensorDataset --- face_swapper/src/data_loader.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/face_swapper/src/data_loader.py b/face_swapper/src/data_loader.py index 1eaadcc..d1145e1 100644 --- a/face_swapper/src/data_loader.py +++ b/face_swapper/src/data_loader.py @@ -5,13 +5,14 @@ from typing import Tuple import cv2 import torch -from torch.utils.data import TensorDataset +from torch import Tensor +from torch.utils.data import Dataset from torchvision import transforms from .types import Batch, ImagePathList, ImagePathSet -class DataLoader(TensorDataset): +class DataLoader(Dataset[Tensor]): def __init__(self, dataset_path : str, dataset_image_pattern : str, dataset_directory_pattern : str, same_person_probability : float) -> None: self.same_person_probability = same_person_probability self.directory_paths = glob.glob(dataset_directory_pattern.format(dataset_path))