Difficulties with implementing DALI on Tensorflow / Keras

Hello, I'm new to this library, and I can't figure out how to properly implement the data preparation process using Tensorflow (Keras). The fact is that I found only the basic examples in the documentation, and in my case, the data is not loaded at once, but is gradually loaded in batches.(like in keras train_generator) It turned out to be difficult to implement, but with PyTorch I did it and it works. I would like to know if there is any alternative to torch DALIClassificationIterator but in dali TF/keras api? And if not, how can I port my implementation to keras? I searched a lot for examples on github, but there are few examples in the public domain, and most of them written in old version of library

class TrainIterator(object):
    def __init__(self, batch_size, device_id, num_gpus):
        self.images_dir = "./dataset/category/data/dataset/"
        self.batch_size = batch_size
        with open(self.images_dir + "train.txt", 'r') as f:
            self.files = [line.rstrip() for line in f if line is not '']

        self.data_set_len = len(self.files)
        self.files = self.files[self.data_set_len * device_id // num_gpus:
                                self.data_set_len * (device_id + 1) // num_gpus]
        self.n = len(self.files)

    def __iter__(self):
        self.i = 0
        shuffle(self.files)
        return self

    def __next__(self):
        batch = []
        labels = []

        if self.i >= self.n:
            self.__iter__()
            raise StopIteration

        for _ in range(self.batch_size):
            jpeg_filename, label = self.files[self.i % self.n].split(' ')
            batch.append(np.fromfile(jpeg_filename, dtype = np.uint8))  # we can use numpy
            labels.append(torch.tensor(int(label), dtype = torch.uint8)) # or PyTorch's native tensors
            self.i += 1
        return (batch, labels)

    def __len__(self):
        return self.data_set_len # 17

    next = __next__

class ValIterator(object):
    def __init__(self, batch_size, device_id, num_gpus):
        self.images_dir = "./dataset/category/data/dataset/"
        self.batch_size = batch_size
        with open(self.images_dir + "val.txt", 'r') as f:
            self.files = [line.rstrip() for line in f if line is not '']
        
        self.data_set_len = len(self.files)
        # based on the device_id and total number of GPUs - world size
        # get proper shard
        self.files = self.files[self.data_set_len * device_id // num_gpus:
                                self.data_set_len * (device_id + 1) // num_gpus]
        self.n = len(self.files)

    def __iter__(self):
        self.i = 0
        shuffle(self.files)
        return self

    def __next__(self):
        batch = []
        labels = []

        if self.i >= self.n:
            self.__iter__()
            raise StopIteration

        for _ in range(self.batch_size):
            jpeg_filename, label = self.files[self.i % self.n].split(' ')
            batch.append(np.fromfile(jpeg_filename, dtype = np.uint8))  # we can use numpy
            labels.append(torch.tensor(int(label), dtype = torch.uint8)) # or PyTorch's native tensors
            self.i += 1
        return (batch, labels)

    def __len__(self):
        return self.data_set_len

    next = __next__

def TrainSourcePipeline(batch_size, num_threads, device_id, external_data):
    pipe = Pipeline(batch_size, num_threads, device_id)
    with pipe:
        jpegs, labels = fn.external_source(source=external_data, num_outputs=2, dtype=types.UINT8)
        images = fn.decoders.image(jpegs, device='mixed')
        images = fn.resize(images, resize_x=224, resize_y=224)
        images = fn.crop_mirror_normalize(images,
                                          dtype=types.FLOAT,
                                          output_layout="CHW",
                                          crop=(224, 224),
                                          mean=[0.485, 0.456, 0.406], 
                                          std=[0.229, 0.224, 0.225],
                                          mirror=False)
        pipe.set_outputs(images, labels)
    return pipe

def ValSourcePipeline(batch_size, num_threads, device_id, external_data):
    pipe = Pipeline(batch_size, num_threads, device_id)
    with pipe:
        jpegs, labels = fn.external_source(source=external_data, num_outputs=2, dtype=types.UINT8)
        images = fn.decoders.image(jpegs, device='mixed')
        images = fn.resize(images, resize_x=224, resize_y=224)
        images = fn.crop_mirror_normalize(images,
                                          dtype=types.FLOAT,
                                          output_layout="CHW",
                                          crop=(224, 224),
                                          mean=[0.485, 0.456, 0.406], 
                                          std=[0.229, 0.224, 0.225],
                                          mirror=False)
        pipe.set_outputs(images, labels)
    return pipe

train_set = TrainIterator(batch_size, 0, 1)
val_set = ValIterator(batch_size, 0, 1)

train_pipe = TrainSourcePipeline(batch_size=batch_size, num_threads=2, device_id = 0, external_data = train_set)
val_pipe = ValSourcePipeline(batch_size=batch_size, num_threads=2, device_id = 0, external_data = val_set)

train_iterator = PyTorchIterator(train_pipe, last_batch_padded=True, last_batch_policy=LastBatchPolicy.PARTIAL, auto_reset=True)
val_iterator = PyTorchIterator(val_pipe, last_batch_padded=True, last_batch_policy=LastBatchPolicy.PARTIAL, auto_reset=True)

Then i just defining criterion, optimizer, model, etc. And starting training loop

for epoch in range(epochs):
    running_loss = 0.0
    for i, data in tqdm(enumerate(train_iterator)):
    ...