@ptrblck I revisited my problem. Your solution is unfortunately not working because logdets is a tensor where I want to let the gradient flow depending on the element-wise condition. i thought about the torch function interface but i am not sure whether it’s working since it provides gradients with respect to the output but I would need to manipulte the gradients flowing for the logdets computation.