Torch.onnx._export fails

Hey there,

I have this model:

class PlayingPolicy(nn.Module):

    def __init__(self):

        # Parameters
        self.gamma = 0.9 # reward discount factor

        # Network
        in_channels = 4 # four players
        out_channels = 8

        self.conv418 = nn.Conv2d(in_channels, out_channels, kernel_size=(3, 4), stride=1, dilation=(1, 8))
        self.conv881 = nn.Conv2d(in_channels, out_channels, kernel_size=(3, 8), stride=8, dilation=(1, 1))
        self.classify = nn.Linear(104, len( len( #  KARO, HERC, PIK, TREF, dalje

        # Optimizer
        self.optimizer = optim.SGD(self.parameters(), lr=1e-2, momentum=0.9)
        # ...

    def forward(self, state: np.ndarray, bidder, trump, legalCards):
        # datatype: <class 'numpy.ndarray'>, <enum 'Suit'>, <class 'int'>, <class 'list'>
        # ....

and when I try to

            input_nn = (playingState, bidderIndex, trumpIndex, [])
            torch.onnx._export(policy, input_nn, path,  export_params=True)

I get this error:

Traceback (most recent call last):
  File "", line 64, in <module>
  File "", line 38, in train_player
  File "/home/markus/Documents/06_Software_Projects/belot/game/", line 402, in saveNetworks
  File "/home/markus/Documents/06_Software_Projects/belot/players/PlayerRL/", line 183, in saveNetwork
    torch.onnx._export(policy, input_nn, path,  export_params=True)
  File "/home/markus/Documents/06_Software_Projects/belot/belot_env/lib/python3.6/site-packages/torch/onnx/", line 26, in _export
    result = utils._export(*args, **kwargs)
  File "/home/markus/Documents/06_Software_Projects/belot/belot_env/lib/python3.6/site-packages/torch/onnx/", line 416, in _export
  File "/home/markus/Documents/06_Software_Projects/belot/belot_env/lib/python3.6/site-packages/torch/onnx/", line 279, in _model_to_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args, training)
  File "/home/markus/Documents/06_Software_Projects/belot/belot_env/lib/python3.6/site-packages/torch/onnx/", line 236, in _trace_and_get_graph_from_model
    trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph(model, args, _force_outplace=True, _return_inputs_states=True)
  File "/home/markus/Documents/06_Software_Projects/belot/belot_env/lib/python3.6/site-packages/torch/jit/", line 277, in _get_trace_graph
    outs = ONNXTracedModule(f, _force_outplace, return_inputs, _return_inputs_states)(*args, **kwargs)
  File "/home/markus/Documents/06_Software_Projects/belot/belot_env/lib/python3.6/site-packages/torch/nn/modules/", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/markus/Documents/06_Software_Projects/belot/belot_env/lib/python3.6/site-packages/torch/jit/", line 332, in forward
    in_vars, in_desc = _flatten(args)
RuntimeError: Only tuples, lists and Variables supported as JIT inputs/outputs. Dictionaries and strings are also accepted but their usage is not recommended. But got unsupported type numpy.ndarray

any ideas how to solve this one?

I alreay checked here but it did not help me…

PS: My Model
(conv418): Conv2d(4, 8, kernel_size=(3, 4), stride=(1, 1), dilation=(1, 8))
(conv881): Conv2d(4, 8, kernel_size=(3, 8), stride=(8, 8))
(classify): Linear(in_features=104, out_features=32, bias=True)
(criterion): PolicyGradientLoss()

state is passed as a numpy array, which is not supported.
Try to pass it as a tensor and check, if it’s working.


that was the solution:

ALL INPUTS and OUTPUTS (of forward) have to be Tensors

input_nn = (torch.tensor(playingState, dtype=torch.float32), torch.tensor(bidderIndex, dtype=torch.int32), torch.tensor(trumpIndex), mask)
            torch.onnx._export(policy, input_nn, path,  export_params=True)
        returned_tensor = torch.zeros(1, 2)
        returned_tensor[:, 0] = action_idx.item()
        returned_tensor[:, 1] = log_action_probability
        return returned_tensor