passing a loss defined as -distr.log_prob(y) as one of the metrics

Hi guys, i was trying to implement the Beta-NLL presented in https://arxiv.org/pdf/2203.09168v2.pdf .

basically the classical negative log-likelihood is multiplied by tf.stop_gradient(estimated_sigma**beta)

i tried with

def NLL2(y, distr): 
    return - tf.stop_gradient(distr.scale )*distr.log_prob(y) 

and it seems it works but the problem is that as the paper says:

"Note that the new loss Lβ−NLL is not meant for performance evaluation, rather it is designed to result
in meaningful gradients. Due to the weighting term, the loss value does not reflect the model’s quality.
The model performance during training should be monitored with the original negative log-likelihood
objective and, optionally, with RMSE for testing the quality of the mean fit"

but when i try to pass the classical NLL

def NLL(y, distr): 
    return -distr.log_prob(y) 

in the metrics

callback = tf.keras.callbacks.EarlyStopping(monitor='loss', patience=300)
optimizer = tf.optimizers.SGD(learning_rate=0.0002,momentum=0.9)
concrete_dropout_BNN = create_concrete_dropout_bnn_model(train_size=train_size)
concrete_dropout_BNN.compile(optimizer=optimizer,
i                  loss=NLL2 , jit_compile=True  ,metrics= [tf.keras.metrics.RootMeanSquaredError(), NLL]
                 ) 

i get the error

AttributeError: 'Tensor' object has no attribute 'log_prob'

Start at the Keras compile call shown in the report and compare the inputs supplied to a metric with the distribution object expected by NLL. Reproduce the Tensor AttributeError while separating the beta-weighted loss from the original NLL metric; done when the original NLL can be monitored without that error.