ptrblck:

self.leach = nn.ModuleList([BasicConv(
            256+512,
            128,
            kernel_size=(1, 1), stride=(1, 1)) for _ in  range(8)])



i followed your way to define the model arch, but also report the trace error when visualizing the model arch…
    # view network architecture
    writer = SummaryWriter('arch', flush_secs=120)
    with writer:
        writer.add_graph(model, torch.rand(32, 3, 250, 250))  # model graph, with input

btw, my model is a multi-task model…
the error:
Traceback (most recent call last):
  File "/home/user1/test.py", line 28, in <module>
    writer.add_graph(model, torch.rand(32, 3, 250, 250))  # model graph, with input
  File "/home/user1/miniconda3/lib/python3.7/site-packages/tensorboardX/writer.py", line 804, in add_graph
    self._get_file_writer().add_graph(graph(model, input_to_model, verbose, profile_with_cuda, **kwargs))
  File "/home/user1/miniconda3/lib/python3.7/site-packages/tensorboardX/pytorch_graph.py", line 324, in graph
    trace = torch.jit.trace(model, args)
  File "/home/user1/miniconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 686, in trace
    traced = _module_class(func, **executor_options)
  File "/home/user1/miniconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1046, in init_then_register
    original_init(self, *args, **kwargs)
  File "/home/user1/miniconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1046, in init_then_register
    original_init(self, *args, **kwargs)
  File "/home/user1/miniconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1482, in __init__
    self._modules[name] = TracedModule(submodule, id_set, optimize=optimize)
  File "/home/user1/miniconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1046, in init_then_register
    original_init(self, *args, **kwargs)
  File "/home/user1/miniconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1482, in __init__
    self._modules[name] = TracedModule(submodule, id_set, optimize=optimize)
  File "/home/user1/miniconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1046, in init_then_register
    original_init(self, *args, **kwargs)
  File "/home/user1/miniconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1482, in __init__
    self._modules[name] = TracedModule(submodule, id_set, optimize=optimize)
  File "/home/user1/miniconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1046, in init_then_register
    original_init(self, *args, **kwargs)
  File "/home/user1/miniconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1469, in __init__
    check_unique(param)
  File "/home/user1/miniconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1461, in check_unique
    raise ValueError("TracedModules don't support parameter sharing between modules")
ValueError: TracedModules don't support parameter sharing between modules

my model arch

class Backbone(nn.Module):
    def __init__(self, num_layers, drop_ratio, mode='ir'):
        super(Backbone, self).__init__()
        assert num_layers in [50, 100, 152], 'num_layers should be 50,100, or 152'
        assert mode in ['ir', 'ir_se'], 'mode should be ir or ir_se'
        blocks = get_blocks(num_layers)
        if mode == 'ir':
            unit_module = bottleneck_IR
        elif mode == 'ir_se':
            unit_module = bottleneck_IR_SE
        self.input_layer = Sequential(Conv2d(3, 64, (3, 3), 1, 1 ,bias=False), 
                                      BatchNorm2d(64), 
                                      PReLU(64))
        self.output_layer = Sequential(BatchNorm2d(512), 
                                       Dropout(drop_ratio),
                                       Flatten(),
                                       Linear(512 * 7 * 7, 512),
                                       BatchNorm1d(512))
        modules = []
        for block in blocks:
            for bottleneck in block:
                modules.append(
                    unit_module(bottleneck.in_channel,
                                bottleneck.depth,
                                bottleneck.stride))
        self.body = Sequential(*modules)

        # for MTL
        self.tower = nn.Sequential(
            nn.Dropout(),
            nn.Linear(512, 32),
            nn.ReLU(),
            nn.Linear(32, 2),
        )

        self.towers = nn.ModuleList([self.tower for _ in range(40)])

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def forward(self,x):
        x = self.input_layer(x)
        x = self.body(x)
        h_shared = self.output_layer(x)
        out = [tower(h_shared) for tower in self.towers]
        return out