Solved by eusouoandre in post #8 
              
              
                Just to let everybody know that’s how I solved. 
The solution, unfortunately, was to  implement my own  simple batched cholesky ( th.cholesky(..., upper=False) ) and then deal with Nan values using  th.isnan . 
import torch as th

# nograd cholesky
def cholesky(A):
    L = th.zeros_like(A)

    for …