FWIW I had implemented something similar before stumbling upon StandardScaler; I gave it a slightly worse name 
class DatasetNorm1d(nn.Module):
    """Records dataset stats. kthxbye"""
    # Similar to: https://scikit-learn.org/stable/modules/generated/sklearn.preprocessing.StandardScaler.html
    def __init__(self, num_features):
        super().__init__()
        self._num_features = num_features
        self.register_buffer('mean', torch.full((self._num_features,), np.nan))
        self.register_buffer('var', torch.full((self._num_features,), np.nan))

    def _is_initialized(self):
        if torch.isnan(self.mean).any() or torch.isnan(self.var).any():
            return False
        return True

    @torch.no_grad()
    def initialize(self, input_batch):
        """
        Args:
            input_batch: Batch that should represent *all* inputs for a given
                dataset.

        Example:

            norm.initialize(torch.cat([x for (x, _) in train_loader]))
        """
        # TODO(eric.cousineau): Use an accurate running computation?
        # See: https://github.com/pytorch/pytorch/blob/480851ad/aten/src/ATen/native/Normalization.cpp#L215-L269
        assert not self._is_initialized()
        N, L = input_batch.shape
        assert L == self._num_features
        assert N > 1
        var_mean = torch.var_mean(input_batch, dim=0)
        self.var.data[:], self.mean.data[:] = var_mean

    def forward(self, x):
        if not self._is_initialized():
            raise RuntimeError("This must be initialized on the dataset!")
        y = (x - self.mean) / torch.sqrt(self.var)
        return y

# Terse test code
class TestStuff(unittest.TestCase):
    def test_dataset_norm(self):
        xs = torch.Tensor([
            [1., 10.],
            [2., 20.],
            [3., 30.],
            [4., 40.],
        ])

        norm = mut.DatasetNorm1d(2)
        with self.assertRaises(RuntimeError):
            norm(xs)

        norm.initialize(xs)
        ys = norm(xs)
        np.testing.assert_array_equal(ys[:, 0].numpy(), ys[:, 1].numpy())
        self.assertEqual(list(norm.state_dict().keys()), ["mean", "var"])