Tensorboardでpytorchモデルを可視化する際のエラー

Question

概要

torchvisionにある、deeplabv3モデルの可視化を試そうとしています。
こちらを参考にして、可視化を試しています。
https://pytorch.org/tutorials/intermediate/tensorboard_tutorial.html

やった事

import torch
from torch import nn
from torchvision import models
from torch.utils.tensorboard import SummaryWriter

deeplabv3 = models.segmentation.deeplabv3_resnet50(pretrained=True)
# default `log_dir` is "runs" - we'll be more specific here
writer = SummaryWriter('runs/deeplabv3')
x = torch.rand(1,3,512,512)
writer.add_graph(deeplabv3, x)

ここで以下のエラーが出ます。

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
~/Documents/4.自学自習/DeepLearning/SemanticSegmentation/torch_test.py in 
----> 293 writer.add_graph(deeplabv3, x)

~/.pyenv/versions/3.7.4/lib/python3.7/site-packages/torch/utils/tensorboard/writer.py in add_graph(self, model, input_to_model, verbose)
    712         if hasattr(model, 'forward'):
    713             # A valid PyTorch model should have a 'forward' method
--> 714             self._get_file_writer().add_graph(graph(model, input_to_model, verbose))
    715         else:
    716             # Caffe2 models do not have the 'forward' method

~/.pyenv/versions/3.7.4/lib/python3.7/site-packages/torch/utils/tensorboard/_pytorch_graph.py in graph(model, args, verbose)
    289             print(e)
    290             print('Error occurs, No graph saved')
--> 291             raise e
    292 
    293     if verbose:

~/.pyenv/versions/3.7.4/lib/python3.7/site-packages/torch/utils/tensorboard/_pytorch_graph.py in graph(model, args, verbose)
    283     with torch.onnx.select_model_mode_for_export(model, torch.onnx.TrainingMode.EVAL):  # TODO: move outside of torch.onnx?
    284         try:
--> 285             trace = torch.jit.trace(model, args)
    286             graph = trace.graph
    287             torch._C._jit_pass_inline(graph)

~/.pyenv/versions/3.7.4/lib/python3.7/site-packages/torch/jit/__init__.py in trace(func, example_inputs, optimize, check_trace, check_inputs, check_tolerance, strict, _force_outplace, _module_class, _compilation_unit)
    953         return trace_module(func, {'forward': example_inputs}, None,
    954                             check_trace, wrap_check_inputs(check_inputs),
--> 955                             check_tolerance, strict, _force_outplace, _module_class)
    956 
    957     if (hasattr(func, '__self__') and isinstance(func.__self__, torch.nn.Module) and

~/.pyenv/versions/3.7.4/lib/python3.7/site-packages/torch/jit/__init__.py in trace_module(mod, inputs, optimize, check_trace, check_inputs, check_tolerance, strict, _force_outplace, _module_class, _compilation_unit)
   1107             func = mod if method_name == "forward" else getattr(mod, method_name)
   1108             example_inputs = make_tuple(example_inputs)
-> 1109             module._c._create_method_from_trace(method_name, func, example_inputs, var_lookup_fn, strict, _force_outplace)
   1110             check_trace_method = module._c._get_method(method_name)
   1111 

RuntimeError: Encountering a dict at the output of the tracer might cause the trace to be incorrect, this is only valid if the container structure does not change based on the module's inputs. Consider using a constant container instead (e.g. for `list`, use a `tuple` instead. for `dict`, use a `NamedTuple` instead). If you absolutely need this and know the side effects, pass strict=False to trace() to allow this behavior.

ちなみに、vgg16などで試したところエラーが出ませんでした。

deeplabv3は複雑な構造をしており、本来ならModuleが入るところにOrderedDictが入ってるからなのかな、という想定ではあります。
しかし、どこをどう変えたらいいのか詳しくわかりません。
こちらについて解決法or知見がある方、教えていただけますと幸いです。