I am trying to implement unet++ for the next day wildfire dataset, but my training process became extremely slow on colab free gpu, I wonder if there is any bug in my code that I am not aware of? Here is the code snipped and also you can find the full code in this link.

Please advise!

def train_model(model: Model, train_dataset: tf.data.Dataset, epochs: int = 10) -> tuple[list[float], list[float]]:
    """
    Trains a model using train dataset. (Save weights of model with best IoU)

    Args:
        model (Model): Model to train.
        train_dataset (Dataset): Training dataset.
        epochs (int): Number of epochs
    Returns:
        Tuple[List[float], List[float]]: Train losses and Validation losses
    """
    optimizer = tf.keras.optimizers.Adam()
    batch_losses = []
    val_losses = []
    best_IoU = 0.0

    for epoch in range(epochs):
        losses = []
        print(f'Epoch {epoch+1}/{epochs}')
        progress = tqdm(train_dataset)
        for images, masks in progress:
            with tf.GradientTape() as tape:
                predictions = model(images, training=True)
                label = tf.where(masks < 0, 0, masks)
                loss = heuristic_loss(label, predictions)
                losses.append(loss.numpy())
                progress.set_postfix({'batch_loss': loss.numpy()})
            gradients = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(gradients, model.trainable_variables))
        print("Evaluation...")
        IoU, recall, precision, val_loss = evaluate_model(lambda x: tf.where(model.predict(x) > 0.5, 1, 0)[:,:,:,0], validation_dataset)
        print("Validation set metrics:")
        print(f"Mean IoU: {IoU}\nMean precision: {precision}\nMean recall: {recall}\nValidation loss: {val_loss}\n")

        print(f'Epoch: {epoch}, Train loss: {np.mean(losses)}')
        batch_losses.append(np.mean(losses))
        val_losses.append(val_loss)

    print(f"Best model IoU: {best_IoU}")
    return batch_losses, val_losses

# Set reproducibility
tf.random.set_seed(1337)

# Create and train the model
segmentation_modelunetPlus = unet_plus_plus(input_shape=(32, 32, 12), num_classes=1)
train_losses, val_losses = train_model(segmentation_modelunetPlus, train_dataset, epochs=15)