Виснет обучение сетки на GPU. Tensorflow

Установил Tensorflow 2.10. CUDA, cuDNN есть, сам тф видит видяху. Проблема в том, что обучение сетки виснет (судя по всему даже не начинается). Как это выглядит в dataspell:

Пытался разобраться, в чем может быть проблема, похоже, что тф все-таки взаимодействует с видеокартой, на графике ГП: первый подскок - подгрузка датасета, второй подскок - инициализация модели, однако при запуске .fit() никакой реакции от ГП не наблюдается.

Грешил на переполнение памяти, но никаких варнингов не было (на всякий случай прикреплю код погрузки данных, буду рад критике):

ds_train_ = image_dataset_from_directory(
    DATA_PATCH_TRAIN,
    labels='inferred',
    label_mode='binary',
    image_size=[128, 128],
    interpolation='nearest',
    batch_size=64,
    shuffle=True,
)
ds_valid_ = image_dataset_from_directory(
    DATA_PATCH_TEST,
    labels='inferred',
    label_mode='binary',
    image_size=[128, 128],
    interpolation='nearest',
    batch_size=64,
    shuffle=False,
)

# Data Pipeline
def convert_to_float(image, label):
    image = tf.image.convert_image_dtype(image, dtype=tf.float32)
    return image, label

AUTOTUNE = tf.data.experimental.AUTOTUNE
ds_train = (
    ds_train_
    .map(convert_to_float)
    .cache()
    .prefetch(buffer_size=AUTOTUNE)
)
ds_valid = (
    ds_valid_
    .map(convert_to_float)
    .cache()
    .prefetch(buffer_size=AUTOTUNE)
)

Честно говоря, не знаю, что еще может пригодиться для решения проблемы, если что, требуйте :) Зарание спасибо