Extending the Inception I3D architecture to TwoStream I3D

[MounirB]

Hello,
I'm willing to extend the pytorch I3D architecture to a TwoStream implementation, like I did in this Keras project, but I'm wondering what changes should be operated to make this architecture behave like the Keras version.
Here are the lines of code I'm stuck onto :

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()

I would like to remove the self.dropout and the self.logits lines. Afterthat, I want to flatten the output of the AvgPool3d and connect it with a Dense layer which does the predictions.
Could you help me in this process, please ?