From 59b4c478be7f9f2570fde50d9dca633f03948eb3 Mon Sep 17 00:00:00 2001
From: Uri Granta <uri.granta@secondmind.ai>
Date: Mon, 9 Sep 2024 10:54:03 +0100
Subject: [PATCH] Check whether register_keras_serializable is actually
 necessary

---
 trieste/models/keras/utils.py | 12 +-----------
 1 file changed, 1 insertion(+), 11 deletions(-)

diff --git a/trieste/models/keras/utils.py b/trieste/models/keras/utils.py
index 741edfdfd6..5fb0bbd972 100644
--- a/trieste/models/keras/utils.py
+++ b/trieste/models/keras/utils.py
@@ -14,23 +14,14 @@
 
 from __future__ import annotations
 
-from typing import Callable, Optional
+from typing import Optional
 
 import tensorflow as tf
 import tensorflow_probability as tfp
-from gpflow.keras import tf_keras
 
 from ...data import Dataset
 from ...types import TensorType
 
-try:
-    register_keras_serializable = tf_keras.saving.register_keras_serializable()
-except AttributeError:  # pragma: no cover (tested but not by coverage)
-
-    # not required in earlier version of TF
-    def register_keras_serializable(func: Callable[..., object]) -> Callable[..., object]:
-        return func
-
 
 def get_tensor_spec_from_data(dataset: Dataset) -> tuple[tf.TensorSpec, tf.TensorSpec]:
     r"""
@@ -131,7 +122,6 @@ def sample_model_index(
     return indices
 
 
-@register_keras_serializable
 def negative_log_likelihood(
     y_true: TensorType, y_pred: tfp.distributions.Distribution
 ) -> TensorType: