Предсказания сохраненной модели отличаются от предсказаний в Юпитере

Предсказания сохраненной модели отличаются от предсказаний в Юпитере. Я сохранял модель и с помощью чекпоинтов, и model.save_model(), во всех случаях предсказания отличаются, при чем сильно в худшую сторону. Немного контекста. Моделька юнетобразная для сегментации, все уже проверял и пришел к выводу, что проблема либо в сохранении модели, либо в ее загрузке. Кастомные метрики не использую. Сама модель: '''

class ImageSegmentationModel:
    def __init__(self, input_shape=(256, 256, 3), num_classes=11, model_path=None, logger=None):
        self.logger = logger if logger else setup_logger(self.__class__.__name__, 'logs/image_segmentation_model.log')
        self.logger.info('Initializing ImageSegmentationModel')

        self.input_shape = input_shape
        self.num_classes = num_classes

        if model_path:
            self.model = self.load_model(model_path)
        else:
            self.model = self._build_model()

    def _build_model(self):
        self.logger.info('Building model...')
        inputs = layers.Input(shape=self.input_shape)
        base_model = applications.EfficientNetB0(include_top=False, weights='imagenet', input_tensor=inputs)
        base_model.trainable = True

        for layer in base_model.layers[:20]:
            layer.trainable = False

        conv1 = base_model.get_layer('block2a_expand_activation').output
        conv2 = base_model.get_layer('block3a_expand_activation').output
        conv3 = base_model.get_layer('block4a_expand_activation').output
        conv4 = base_model.get_layer('block6a_expand_activation').output
        conv5 = base_model.get_layer('top_activation').output

        up6 = layers.Conv2DTranspose(512, 2, strides=(2, 2), padding='same')(conv5)
        up6 = layers.concatenate([up6, conv4], axis=3)
        conv6 = layers.Conv2D(512, 3, activation='relu', padding='same')(up6)
        conv6 = layers.Conv2D(512, 3, activation='relu', padding='same')(conv6)

        up7 = layers.Conv2DTranspose(256, 2, strides=(2, 2), padding='same')(conv6)
        up7 = layers.concatenate([up7, conv3], axis=3)
        conv7 = layers.Conv2D(256, 3, activation='relu', padding='same')(up7)
        conv7 = layers.Conv2D(256, 3, activation='relu', padding='same')(conv7)

        up8 = layers.Conv2DTranspose(128, 2, strides=(2, 2), padding='same')(conv7)
        up8 = layers.concatenate([up8, conv2], axis=3)
        conv8 = layers.Conv2D(128, 3, activation='relu', padding='same')(up8)
        conv8 = layers.Conv2D(128, 3, activation='relu', padding='same')(conv8)

        up9 = layers.Conv2DTranspose(64, 2, strides=(2, 2), padding='same')(conv8)
        up9 = layers.concatenate([up9, conv1], axis=3)
        conv9 = layers.Conv2D(64, 3, activation='relu', padding='same')(up9)
        conv9 = layers.Conv2D(64, 3, activation='relu', padding='same')(conv9)

        up10 = layers.Conv2DTranspose(32, 2, strides=(2, 2), padding='same')(conv9)
        conv10 = layers.Conv2D(32, 3, activation='relu', padding='same')(up10)
        conv10 = layers.Conv2D(32, 3, activation='relu', padding='same')(conv10)

        outputs = layers.Conv2D(self.num_classes, 1, activation='softmax')(conv10)

        model = models.Model(inputs=[inputs], outputs=[outputs])

        self.logger.info('Model built successfully.')
        return model

'''

Скрипт тренировки и тут же сохранение: '''

import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import tensorflow_datasets as tfds
from src.image_segmentation_model import ImageSegmentationModel
from src.logger import setup_logger

def main():
    logger = setup_logger('train_script', 'logs/train.log')
    logger.info('Starting training script...')

    model = ImageSegmentationModel()

    dataset = model.load_dataset()
    images, masks = model.prepare_data(dataset)

    train_images, train_masks = images[:2000], masks[:2000]
    val_images, val_masks = images[4400:], masks[4400:]

    model.train(train_images, train_masks, val_images, val_masks)

    model.save('model1.keras')

    logger.info('Training script completed.')

if __name__ == "__main__":
    main()

'''

Тест предсказаний: '''

import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import cv2
import numpy as np
from src.image_segmentation_model import ImageSegmentationModel
from src.logger import setup_logger
import matplotlib.pyplot as plt
import tensorflow as tf

def main(image_path):
    logger = setup_logger('predict_script', 'logs/predict.log')
    logger.info(f'Starting prediction script for image: {image_path}')
    
    model_path = 'model1.keras'
    segmentation_model = ImageSegmentationModel(model_path=model_path, logger=logger)
    segmentation_model.model.summary()
    image = cv2.imread(image_path)
    if image is None:
        logger.error(f'Failed to load image at {image_path}')
        sys.exit(1)
    
    image = cv2.resize(image, (256, 256))
    image = image / 255.0
    image = image.astype(np.float32)
    image = np.expand_dims(image, axis=0)  # Добавляем размер батча

    prediction = segmentation_model.model.predict(image)
    prediction_mask = np.argmax(prediction, axis=-1)[0]  # Извлекаем маску из батча

    visualize_prediction(image[0], prediction_mask)
    logger.info('Prediction script completed.')

def visualize_prediction(image, prediction):
    plt.figure(figsize=(10, 10))

    plt.subplot(1, 2, 1)
    plt.title("Input Image")
    plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

    plt.subplot(1, 2, 2)
    plt.title("Predicted Mask")
    plt.imshow(prediction, cmap='jet', alpha=0.5)

    plt.show()

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python predict.py <path_to_image>")
        sys.exit(1)

    image_path = sys.argv[1]
    main(image_path)

'''

Буду очень благодарен.