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):
super().__init__()
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):
pass
def forward(inputs):
return self.layer(inputs)
class TestModule2(nn.Module):
def __init__(self):
super().__init__()
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)))
print(pretty_string)
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!