Using keras model with torch.func.jacrev gives zero gradients

Using keras model with torch.func.jacrev gives zero gradients #1781