How to apply a function to network output before passing it to the loss?

I'm trying to implement a network in tensorflow and I need to apply a function f to the network output and use the returned value as the prediction to be used in the loss.

Is there a simple way to make it or which part of tensorflow should I study to achieve that ?

you should study how to write custom training loops in tensorflow: https://www.tensorflow.org/guide/keras/writing_a_training_loop_from_scratch

A simplified and short version could look similar to the code bellow:

#Repeat for several epochs
for epoch in range(epochs):

  # Iterate over the batches of the dataset.
  for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):

      # Start tracing your forward pass to calculate gradients
      with tf.GradientTape() as tape:
          prediction = model(x_batch_train, training=True)
        
          # HERE YOU PLACE YOUR FUNCTION f
          transformed_prediction = f(prediction)

          loss_value = loss_fn(y_batch_train, transformed_prediction )
      grads = tape.gradient(loss_value, model.trainable_weights)
      optimizer.apply_gradients(zip(grads, model.trainable_weights))
    (...)