Repository: pytorch
Issue: [TensorBoard] Graph with objects other than torch.nn.Module can not be visualized.
Labels: triaged,module: tensorboard

## 🐛 Bug

I use tensorboardX to `add_graph(model, (dump,))`, but 
```
    writer_dict['writer'].add_graph(model, (dump_input, ))
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorboardX/writer.py", line 774, in add_graph
    self._get_file_writer().add_graph(graph(model, input_to_model, verbose, **kwargs))
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorboardX/pytorch_graph.py", line 275, in graph
    trace = torch.jit.trace(model, args)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 772, in trace
    check_tolerance, _force_outplace, _module_class)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 898, in trace_module
    module = make_module(mod, _module_class, _compilation_unit)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 669, in make_module
    return _module_class(mod, _compilation_unit=_compilation_unit)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1386, in init_then_register
    original_init(self, *args, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1386, in init_then_register
    original_init(self, *args, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1881, in __init__
    self._modules[name] = TracedModule(submodule, id_set)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1386, in init_then_register
    original_init(self, *args, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1881, in __init__
    self._modules[name] = TracedModule(submodule, id_set)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1386, in init_then_register
    original_init(self, *args, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1881, in __init__
    self._modules[name] = TracedModule(submodule, id_set)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1386, in init_then_register
    original_init(self, *args, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1881, in __init__
    self._modules[name] = TracedModule(submodule, id_set)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1386, in init_then_register
    original_init(self, *args, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1881, in __init__
    self._modules[name] = TracedModule(submodule, id_set)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1386, in init_then_register
    original_init(self, *args, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1855, in __init__
    assert(isinstance(orig, torch.nn.Module))
AssertionError
```

## To Reproduce

Steps to reproduce the behavior:

```python
import torch
from tensorboardX import SummaryWriter

writer = SummaryWriter()
model = mymodel()
dump_input = torch.rand( (1, 3,256, 256)  )
writer.add_graph(model, (dump_input, ))
```

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

## Environment
Collecting environment information...
PyTorch version: 1.2.0
Is debug build: No
CUDA used to build PyTorch: 10.0.130

OS: Ubuntu 16.04.6 LTS
GCC version: (Ubuntu 5.4.0-6ubuntu1~16.04.11) 5.4.0 20160609
CMake version: Could not collect

Python version: 3.7
Is CUDA available: Yes
CUDA runtime version: 10.1.105
GPU models and configuration: 
GPU 0: GeForce RTX 2080 Ti
GPU 1: GeForce RTX 2080 Ti

Nvidia driver version: 418.43
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.7.5.0
/usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudnn.so.7

Versions of relevant libraries:
[pip3] numpy==1.16.2
[pip3] numpydoc==0.8.0
[pip3] torch==1.2.0
[pip3] torchvision==0.4.0
[conda] blas                      1.0                         mkl  
[conda] mkl                       2019.1                      144  
[conda] mkl-service               1.1.2            py37he904b0f_5  
[conda] mkl_fft                   1.0.10           py37ha843d7b_0  
[conda] mkl_random                1.0.2            py37hd81dba3_0  
[conda] torch                     1.2.0                    pypi_0    pypi
[conda] torchvision               0.3.0                    pypi_0    pypi

## Additional context

The same problem occurred when using `torch.utils.tensorboard`