Using AffineCoupling on all dimensions

Is it possible to use the transforms.AffineCoupling, which takes in a DenseNN and apply the location-scale transformation on all components of the previous random variable.

E.g., to declare a 2d Gaussian with a learnable mean, I would think that one would do

split_dim = 0
dim = 2
hypernet = DenseNN(split_dim, [10*dim], [dim-split_dim, dim-split_dim])
affine = distributions.transforms.AffineCoupling(split_dim,hypernet)
td = distributions.TransformedDistribution(distributions.Normal(torch.zeros(dim),torch.ones(dim)), [affine])

However, this only works when split_dim >= 1. Am I missing something, or is there another way to do this (except passing a nn.Linear and re-defining the base distribution every time)?