Triplet loss causes NotImplementedError

[aoeu1234]

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Fedora 32
• TensorFlow version and how it was installed (source or binary): 2.4.1 (from pip)
• TensorFlow-Addons version and how it was installed (source or binary): 0.12.1 (from pip)
• Python version: Python 3.8.7
• Is GPU used? (yes/no): no

Describe the bug
Using any triplet loss causes NotImplementedError error.

Code to reproduce the issue
Using the example from the docs:

y_true = tf.convert_to_tensor([0, 0])
y_pred = tf.convert_to_tensor([[0.0, 1.0], [1.0, 0.0]])

for f in [tfa.losses.triplet_hard_loss, tfa.losses.triplet_semihard_loss]:
    try:
        f(y_true, y_pred, distance_metric="L2")
    except NotImplementedError:
        print("Error")

for f in [tfa.losses.TripletHardLoss(), tfa.losses.TripletSemiHardLoss()]:
    try:
        f(y_true, y_pred)
    except NotImplementedError:
        print("Error")

Output:

Error
Error
Error
Error

Other info / logs

Actual traceback of the errors (the same, except for function name):

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/zunny/.local/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 828, in __call__
    result = self._call(*args, **kwds)
  File "/home/zunny/.local/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 862, in _call
    results = self._stateful_fn(*args, **kwds)
  File "/home/zunny/.local/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 2941, in __call__
    filtered_flat_args) = self._maybe_define_function(args, kwargs)
  File "/home/zunny/.local/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 3361, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/home/zunny/.local/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 3196, in _create_graph_function
    func_graph_module.func_graph_from_py_func(
  File "/home/zunny/.local/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py", line 990, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/home/zunny/.local/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 634, in wrapped_fn
    out = weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/home/zunny/.local/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py", line 977, in wrapper
    raise e.ag_error_metadata.to_exception(e)
NotImplementedError: in user code:

    /home/zunny/.local/lib/python3.8/site-packages/tensorflow_addons/losses/triplet.py:261 triplet_hard_loss  *
        pdist_matrix = metric_learning.pairwise_distance(
    /home/zunny/.local/lib/python3.8/site-packages/tensorflow_addons/losses/metric_learning.py:66 pairwise_distance  *
        mask_offdiagonals = tf.ones_like(pairwise_distances) - tf.linalg.diag(
    /home/zunny/.local/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py:201 wrapper  **
        return target(*args, **kwargs)
    /home/zunny/.local/lib/python3.8/site-packages/tensorflow/python/ops/array_ops.py:3120 ones
        output = _constant_if_small(one, shape, dtype, name)
    /home/zunny/.local/lib/python3.8/site-packages/tensorflow/python/ops/array_ops.py:2804 _constant_if_small
        if np.prod(shape) < 1000:
    <__array_function__ internals>:5 prod
        
    /home/zunny/.local/lib/python3.8/site-packages/numpy/core/fromnumeric.py:3030 prod
        return _wrapreduction(a, np.multiply, 'prod', axis, dtype, out,
    /home/zunny/.local/lib/python3.8/site-packages/numpy/core/fromnumeric.py:87 _wrapreduction
        return ufunc.reduce(obj, axis, dtype, out, **passkwargs)
    /home/zunny/.local/lib/python3.8/site-packages/tensorflow/python/framework/ops.py:852 __array__
        raise NotImplementedError(

    NotImplementedError: Cannot convert a symbolic Tensor (strided_slice:0) to a numpy array. This error may indicate that you're trying to pass a Tensor to a NumPy call, which is not supported