src/transformers/models/bit/modeling_bit.py/WeightStandardizedConv2d/forward
class WeightStandardizedConv2d:    def forward(self, hidden_state):
        if self.pad is not None:
            hidden_state = self.pad(hidden_state)
        weight = nn.functional.batch_norm(
            self.weight.reshape(1, self.out_channels, -1), None, None, training=True, momentum=0.0, eps=self.eps
        ).reshape_as(self.weight)
        hidden_state = nn.functional.conv2d(
            hidden_state, weight, self.bias, self.stride, self.padding, self.dilation, self.groups
        )
        return hidden_state