diff --git a/declearn/model/tensorflow/utils/_loss.py b/declearn/model/tensorflow/utils/_loss.py
index c6b880c8e8aaa6dc90cb197e18bf7677f32285fc..d70e26a0fe489e7ef95b37b22f8cdcd0e6ffde30 100644
--- a/declearn/model/tensorflow/utils/_loss.py
+++ b/declearn/model/tensorflow/utils/_loss.py
@@ -96,7 +96,7 @@ def build_keras_loss(
                 loss = tf.keras.losses.deserialize(loss)
                 loss.reduction = reduction
             # If this failed, try wrapping the function using LossFunction.
-            except ValueError:
+            except (AttributeError, ValueError):
                 loss = LossFunction(cls)
         # Case when the string was deserialized into a class.
         else: