Question about Resnet's Bottleneck incoming object

I’m optimizing Resnet50, and besides input x, I want to input another objective at each layer.I’m optimizing Resnet50, and besides input x, I want to input another objective at each layer. That is, I want to pass parameters in Bottleneck’s forward function.Every layer of Resnet in many example codes is built by make_layer, which makes it impossible for me to pass other parameters.

class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(Bottleneck, self).__init__()
        ...........

    def forward(self, x, mask):
        residual = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)
        if self.downsample is not None:
            residual = self.downsample(x)
        out += residual
        out = self.relu(out)

        return out

All in all, I just want to know how to pass in the mask in Bottleneck.

You could create a custom nn.Module which accepts the mask object and apply it in the forward.
If you keep the module names equal and only modify the forward method loading the pre-trained state_dict should still work.