I’ve tried manually merging batchnorm and conv layers. A sample of code:
def merge_bn_conv(bn_layer, conv_layer):
    weight = (bn_layer.weight.detach()*conv_layer.weight.permute(1,2,3,0)/(1e-10 + bn_layer.running_var.detach().sqrt())).permute(3,0,1,2)
    conv_layer.weight = nn.Parameter(weight)
    bias = bn_layer.bias.detach() - bn_layer.running_mean.detach()*bn_layer.weight.detach()/(1e-10 + bn_layer.running_var.detach().sqrt())
    conv_layer.bias = nn.Parameter(bias)

In modules:
def drop_batchnorm():
    merge_bn_conv(self.bn1, self.conv1)
    del self.bn1
    ...
    self.other_module.drop_batchnorm()
    ...
    self.batchnorm_dropped = True

In .forward:
def forward(self, input):
    ...
    out = self.conv1(out)
    if not self.batchnorm_dropped:
        out = self.bn1(out)
    ...

I compared the same model in eval mode before and after calling drop_batchnorm.
Surprisingly, after drop_batchnorm the throughput of the model dropped by around 10%…
I could understand if there was no difference (pytorch doing this optimization for you), but how can this optimization actually make things worse?
Would be cool if someone clarified this (and double checked that the effect exists).