INNER CODE UNIT · Python
WrappedDataset
TIGER-AI-Lab/ImagenHub · src/imagen_hub/pipelines/custom_diffusion/custom_diffusion_src/stable-diffusion/main.py:136
class WrappedDataset(Dataset):
"""Wraps an arbitrary object with __len__ and __getitem__ into a pytorch dataset"""
def __init__(self, dataset):
self.data = dataset
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
return self.data[idx]
def worker_init_fn(_):
worker_info = torch.utils.data.get_worker_info()
dataset = worker_info.dataset
worker_id = worker_info.id