Problems using tf.linalg.det as loss function

I'm training a model with components that are represented by square matrices. The closer the determinant of those matrices is to 1 the better, so I implemented a loss function as follows :

@tf.function
def det_loss(slice):
  out = tf.linalg.det(slice)
  return out

(I left subtracting the absolute or square value from 1 out, this is enough to reproduce my error)

My train_step function looks as follows:

@tf.function
def train_step(x, seed, pad_target):
    iter_n = tf.random.uniform([], 169, 196, tf.int32)
    with tf.GradientTape() as g:
        for i in tf.range(iter_n):
            x = ca(x,seed=seed)

        out_loss, target_loss = loss_f(x,seed,pad_target)
        target_loss = tf.reduce_mean(target_loss)
        loss = tf.reduce_mean(out_loss)

    grads = g.gradient(loss, ca.weights)
    grads = [g / (tf.norm(g) + 1e-8) for g in grads]
    trainer.apply_gradients(zip(grads, ca.weights))
    return x, loss, target_loss

loss_f takes the input and slices it into 6x6 matrices on which the det_loss is called. The problem is not with the slicing as only the det_loss does not work with it.

Now when I run the code the foward pass works fine but when the gradient is calculated I get the following error:

Traceback (most recent call last):
  File "C:/Users/.../main.py", line 207, in <module>
    x, loss, target_loss = train_step(x0,seed, pad_target)
  File "C:\Users\...\venv\lib\site-packages\tensorflow\python\util\traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "C:\Users\...\venv\lib\site-packages\tensorflow\python\eager\execute.py", line 58, in quick_execute
    tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
tensorflow.python.framework.errors_impl.InvalidArgumentError: 2 root error(s) found.
  (0) INVALID_ARGUMENT:  Input is not invertible.
     [[node gradients/MatrixDeterminant_grad/MatrixInverse
 (defined at C:/Users/.../main.py:176)
]]
     [[StatefulPartitionedCall/StatefulPartitionedCall/StatefulPartitionedCall_749/MatrixDeterminant/_4508]]
  (1) INVALID_ARGUMENT:  Input is not invertible.
     [[node gradients/MatrixDeterminant_grad/MatrixInverse
 (defined at C:/Users/.../main.py:176)
]]
0 successful operations.
0 derived errors ignored. [Op:__inference_train_step_65237]

Errors may have originated from an input operation.
Input Source operations connected to node gradients/MatrixDeterminant_grad/MatrixInverse:
In[0] gradients/MatrixDeterminant_grad/MatrixInverse/slice:

Operation defined at: (most recent call last)
>>>   File "C:/Users/.../main.py", line 207, in <module>
>>>     x, loss, target_loss = train_step(x0,seed, pad_target)
>>> 
>>>   File "C:/Users/.../main.py", line 176, in train_step
>>>     grads = g.gradient(loss, ca.weights)
>>> 

Input Source operations connected to node gradients/MatrixDeterminant_grad/MatrixInverse:
In[0] gradients/MatrixDeterminant_grad/MatrixInverse/slice:

Operation defined at: (most recent call last)
>>>   File "C:/Users/.../main.py", line 207, in <module>
>>>     x, loss, target_loss = train_step(x0,seed, pad_target)
>>> 
>>>   File "C:/Users/.../main.py", line 176, in train_step
>>>     grads = g.gradient(loss, ca.weights)
>>> 

Function call stack:
train_step -> __backward_det_loss_13020 -> train_step -> __backward_det_loss_13020


Process finished with exit code 1

This looks to me like tf has problems calculating the gradient for inverting a matrix. I have tried implementing a custom gradient to no avail.

This might be completely wrong but I'm also confused why there needs to be a gradient for this in the first place as it is not part of the model itself if that makes sense.

I'm using Python 3.8.0 and tensorflow 2.7.0