cond.py

import torch
import functorch
from functorch import vmap, grad

def foo(input, p):
    result = torch.linalg.cond(input, p)
    return result.sum()

input = torch.randn(5, 5)
p = 1

grad(foo)(input, p)