I think BatchNorm1d might not .detach() the mean and variance when normalizing. When I use the implementation below it seems work (simplified by removing the extra bias and weight).
class BatchNorm(nn.Module):
    def __init__(self, nFeatures, eps=1e-5, momentum=0.1):
        super().__init__()
        self.register_buffer("moving_avg", torch.zeros(nFeatures))
        self.register_buffer("moving_var", torch.ones(nFeatures))
        self.register_buffer("eps", torch.tensor(eps))
        self.register_buffer("momentum", torch.tensor(momentum))
    
    def forward(self, x):        
        if self.training:
            mean = x.mean(dim=0)
            var = x.var(dim=0)
            self.moving_avg = self.moving_avg * self.momentum + mean * (1 - self.momentum)
            self.moving_var = self.moving_var * self.momentum + var * (1 - self.momentum)

        else:
            mean = self.moving_avg
            var = self.moving_var
            
        # don't think original implementation uses .detach() on mean & var
        x_norm = (x - mean.detach()) / (torch.sqrt(var.detach() + self.eps))
        return x_norm

I tried to find the exact implementation used by PyTorch to check. In torch.nn.functional the batch_norm function calls torch.batch_norm. Could somebody point me into the direction of where that’s implemented?