fix import

This commit is contained in:
rromb 2022-05-31 15:56:29 +02:00
parent 174340dd4c
commit fdc93bd76d

View file

@ -18,7 +18,7 @@ from webdataset.handlers import warn_and_continue
from ldm.util import instantiate_from_config from ldm.util import instantiate_from_config
class DataWithWings(data.IterableDataset): class DataWithWings(torch.utils.data.IterableDataset):
def __init__(self, min_size, transform=None, target_transform=None): def __init__(self, min_size, transform=None, target_transform=None):
self.min_size = min_size self.min_size = min_size
self.transform = transform if transform is not None else nn.Identity() self.transform = transform if transform is not None else nn.Identity()