Tensorboard add_graph (actually jit.trace) does not work with None inputs

tensorboard
1

Hi, I have a problem where a forward method takes optional inputs:

def forward(self, x_in):
    x1, x2 = x_in
   if x2 is None:
     ...

so when I want to add_graph for a case without the x2 provided…

writer.add_graph(model, input_to_model=((x, None)))

… I get an error:

"**/jit/_trace.py", line 1278, in trace_module
    module._c._create_method_from_trace(
RuntimeError: Type 'Tuple[Tensor, NoneType]' cannot be traced. Only Tensors and (possibly nested) Lists, Dicts, and Tuples of Tensors can be traced

Is there an easy workaround?

  • changing the signature of forward is not an option
  • in this simple example, also providing x2 would not hurt, but the real example is more complex, resulting in another graph

Thanks in advance

(Minimal reproducible example on request available)

2

btw: I found out there is an open request for something similar, but its open since a while… Allow tracing of models which output `None` · Issue #18850 · pytorch/pytorch · GitHub

the only solution I came up with, is either wrapping or monkey-patching the forward method of the model to not take a tuple but optional arguments…