We can use NFs not for variance estimation only, but for covariance estimation as well. For this we need to sample input batch with the same random seeds. I've found that it is probably achievable with torch.func.vmap(randomness='same') stabilized in PyTorch v2 (see #7 for torch upgrading issue), but I cannot make it work.