Could not export Python function call 'quantum_circuit'. Remove calls to Python functions before export

I can’t save the model, using:

    traced_cell = torch.jit.script(model)
    # saving model to disc memory for post processing
    traced_cell.save(path + f"/policy_{sample}.pt")

The error:

"Could not export Python function call 'quantum_circuit'. Remove calls to Python functions before export. Did you forget to add @script or @script_method annotation? If this is a nn.ModuleList, add it to __constants__:"

Could any one can help me?
Note: I use the default qubit mode:

dev = qml.device("default.qubit", wires=n_qubits); 

@qml.qnode(dev)
def quantum_circuit(inputs, weights):
    weights = weights.reshape(q_depth, n_qubits)
    # Initialise latent vectors
    for i in range(n_qubits):
        qml.RY(inputs[i], wires=i)
    # Repeated layer
    for i in range(q_depth):
        # Parameterised layer
        for y in range(n_qubits):
            qml.RY(weights[i][y], wires=y)
        # Control Z gates
        for y in range(n_qubits - 1):
            qml.CZ(wires=[y, y + 1])
    return qml.probs(wires=list(range(n_qubits)))

Hi @Hongfu_Zhang, welcome to the Forum!

This sound like a Torch question more than a PennyLane question. Have you tried using torch.save and torch.load?

If you think it’s an issue with PennyLane and need more help, can you please share a minimal reproducible example that we can use to try to replicate your error?

A minimal reproducible example (or minimal working example) is the simplest version of the code that reproduces the problem. It should be self-contained, including all necessary imports, data, functions, etc., so that we can copy-paste the code and reproduce the problem. However it shouldn’t contain any unnecessary data, functions, etc. (for example gates and functions that can be removed to simplify the code). If you’re not sure what this means then please make sure to check out this video.