Hello @ptrblck,
Thank you for your answer.
My purpose is to load a pretrained model and then modify the architecture of the network by pruning its last layers.
end_point = 'Logits'
self.avg_pool = nn.AvgPool3d(kernel_size=[2, 7, 7],
                                     stride=(1, 1, 1))
# self.dropout = nn.Dropout(dropout_keep_prob)
self.logits = Unit3D(in_channels=384+384+128+128, output_channels=self._num_classes,
                             kernel_shape=[1, 1, 1],
                             padding=0,
                             activation_fn=None,
                             use_batch_norm=False,
                             use_bias=True,
                             name='logits')

self.build()

After a pretrained model is loaded with the following
i3d = InceptionI3d(400, in_channels=2)
i3d.load_state_dict(torch.load(os.path.join(save_model, 'flow_imagenet.pt')))

I try to modify the structure of the architecture to remove Unit3D and replace it by a Dense layer
def replace_logitsWithLinear(self, num_classes):
        self._num_classes = num_classes
        self.logits = nn.Linear(4096, self._num_classes)
        self.add_module('logits', self.logits)
        print(list(self.children())[:-1])

However, I get the following error :
RuntimeError: size mismatch, m1: [4096 x 1], m2: [4096 x 11] at /pytorch/aten/src/THC/generic/THCTensorMathBlas.cu:283

I guess it’s a mismatch between the size of the AvgPool3d and the Linear layer.