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))