ONNX export if submodule has extra state

How may I export a module with a submodule that implements get_extra_state() but doesn’t return a tensor (returns a dict maybe)?
For example, when I run the following code:

from typing import Any

import torch
from torch import nn

class TestModule1(nn.Module):
    def __init__(self):
        self.layer = nn.Linear(10, 10)

    def get_extra_state(self) -> Any:
        return {"value1": 42, "value2": "very important"}

    def set_extra_state(self, state: Any):

    def forward(inputs):
        return self.layer(inputs)

class TestModule2(nn.Module):
    def __init__(self):
        self.test = TestModule1()
        self.relu = nn.ReLU()

    def forward(inputs):
        x = self.test(inputs)
        return self.relu(x)

module = TestModule2()
pretty_string = torch.onnx.export_to_pretty_string(module, torch.ones((1, 10)))

it fails with the error:

  File ".../lib/python3.10/site-packages/torch/onnx/__init__.py", line 388, in export_to_pretty_string
    return utils.export_to_pretty_string(*args, **kwargs)
  File ".../lib/python3.10/site-packages/torch/onnx/utils.py", line 852, in export_to_pretty_string
    graph, params_dict, torch_out = _model_to_graph(
  File ".../lib/python3.10/site-packages/torch/onnx/utils.py", line 727, in _model_to_graph
    graph, params, torch_out, module = _create_jit_graph(model, args)
  File ".../lib/python3.10/site-packages/torch/onnx/utils.py", line 602, in _create_jit_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args)
  File ".../lib/python3.10/site-packages/torch/onnx/utils.py", line 515, in _trace_and_get_graph_from_model
    orig_state_dict_keys = torch.jit._unique_state_dict(model).keys()
  File ".../lib/python3.10/site-packages/torch/jit/_trace.py", line 71, in _unique_state_dict
    filtered_dict[k] = v.detach()
AttributeError: 'dict' object has no attribute 'detach'

I would appreciate any help!