I tried to use torch.jit.trace and it was able to successfully trace the model.
traced_module = torch.jit.trace(model, input_tensor) # success

But I still get the same error when I try to use the traced module in the onnx export.
    torch.onnx.export(
        traced_module,
        input_tensor,
        onnx_file_path,
        input_names=["input"],
        output_names=["output"],
        export_params=True,
    ) # fail, same as before