You should be able to use torch.nn.init, if you check for None before calling an init method on the affine parameters:
def bn_weight_init(m):
    if m.weight is not None and m.bias is not None:
        torch.nn.init.normal_(m.weight)
        torch.nn.init.normal_(m.bias)

bn = nn.BatchNorm2d(3, affine=True)
print(bn.weight, bn.bias)
bn.apply(bn_weight_init)
print(bn.weight, bn.bias)

bn = nn.BatchNorm2d(3, affine=False)
print(bn.weight, bn.bias)
bn.apply(bn_weight_init)
print(bn.weight, bn.bias)