Vectorizing a normalizing flow

Hello, I am quite new to Pytorch and DL in general. I have a flow (an NVP one to be exact) that takes a 2d tensor z in input and transforms it into another 2d tensor \phi, like so

phi,logJ = the_Flow(z)

and returns also the logarithm of the Jacobian of the transformation.
I would like to be able to feed a batch of zs (i.e. a 3d tensor) without having to do it sequentially (in a “for” loop) for every single z, and get a batch of \phis and determinants in return. I need to do it because the loss needs to be computed as an average over said batch.
Can I do this without having to redefine the flow to operate on 3d tensors?