Ошибка: Lambda fuctions will be no more assumed to be used in the statement where they are used

делаю нейронку и пытаюсь нормализировать данные с помощью вот такой конструкции:

normalized_ds = train_ds.map(lambda x,y: (normalization_layer(x), y))
image_batch, labels_batch = next(iter(normalized_ds))
first_image = image_batch[0]```

но происходит вот такая ошибка:
Lambda fuctions will be no more assumed to be used in the statement where they are used, or at least in the same block.