TorchScripted loss function breaks in the backward

I have a loss function which is torchscripted using torch.jit.script decorator. Before torchscripting it works fine but after torchscripting it fails in the backward path due to some in-place operation. However, I cannot find any in-place operation in my loss function. Any idea?