mirror of
https://github.com/HumanAIGC/lite-avatar.git
synced 2026-02-05 01:49:19 +08:00
23 lines
453 B
Python
23 lines
453 B
Python
from torch.utils.data import IterableDataset
|
|
|
|
|
|
def default_fn(data):
|
|
return data
|
|
|
|
|
|
class MapperIterDataPipe(IterableDataset):
|
|
|
|
def __init__(self,
|
|
datapipe,
|
|
fn=default_fn):
|
|
self.datapipe = datapipe
|
|
self.fn = fn
|
|
|
|
def set_epoch(self, epoch):
|
|
self.epoch = epoch
|
|
|
|
def __iter__(self):
|
|
assert callable(self.fn)
|
|
for data in self.datapipe:
|
|
yield self.fn(data)
|