I'm trying to build a custom loss function which utilizes the gradient of the outputs predicted by the model (i.e. gradient of y_pred). For example, the loss function may try to minimize the mean squared error between the gradient of the model output (grad(y_pred)) and y_true.
Although many operations like taking the square of y_pred inside the custom loss function is okay, using its gradient within the custom loss function is lnot possible. I understand that this might be because y_pred is simply a reference to a computational graph, and doesn't have any values initially. Can someone please suggest a workaround for this?