`example_input_array` with `Task` like usage of LightningModule

How can I get example_input_array work with Task like usage of Lightning as mentioned here. Currently, I get RuntimeError: The traced function didn't return any values! Side-effects are not captured in traces, so it would be a no-op.

Setup:

  • Trainer flag weights_summary="full".
  • Using a LightningModule as Task. i.e. I don’t have the forward method.
  • LightningModule has example_input_array attribute defined.

Here is the trace

File "/home/zoro/anaconda3/lib/python3.8/site-packages/pytorch_lightning/loggers/tensorboard.py", line 190, in log_graph
    self.experiment.add_graph(
  File "/home/zoro/anaconda3/lib/python3.8/site-packages/torch/utils/tensorboard/writer.py", line 714, in add_graph
    self._get_file_writer().add_graph(graph(model, input_to_model, verbose))
  File "/home/zoro/anaconda3/lib/python3.8/site-packages/torch/utils/tensorboard/_pytorch_graph.py", line 291, in graph
    raise e
  File "/home/zoro/anaconda3/lib/python3.8/site-packages/torch/utils/tensorboard/_pytorch_graph.py", line 285, in graph
    trace = torch.jit.trace(model, args)
  File "/home/zoro/anaconda3/lib/python3.8/site-packages/torch/jit/__init__.py", line 953, in trace
    return trace_module(func, {'forward': example_inputs}, None,
  File "/home/zoro/anaconda3/lib/python3.8/site-packages/torch/jit/__init__.py", line 1109, in trace_module
    module._c._create_method_from_trace(method_name, func, example_inputs, var_lookup_fn, strict, _force_outplace)
RuntimeError: The traced function didn't return any values! Side-effects are not captured in traces, so it would be a no-op.

This should be solved now in the latest version 0.9.1rc1.
0.9.0 unfortunately had log_graph=True by default, which was a major oversight because not all models implement forward, and so logging a graph only makes sense if you have a forward. Since Task does not have that, it makes no sense to log the graph.

If you need to stick with 0.9.0, set log_graph = False in your Logger, otherwise I recommend you upgrade to 0.9.1rc1