Running keras model.predict with generator (or model.predict_generator) automatically converts to tensor

Hello, I running a model which is just Keras' built-in Xception run on a large image dataset but with TFP distribution layers as output. Because I am using images, I have to use generators to train/test as loading all of my train/test data into memory is not possible.

I've noticed that when using model.predict(testgenerator) with a generator, rather than returning distribution objects as might be obtained when running model(testdata), a numpy tensor is expected and therefore the convert_to_tensor function is automatically run. Thus, to get all of the distribution information I want (for example, mean and covariance matrix for a multivariate gaussian), I have to run the prediction twice, once with convert to tensor set to mean, and once for covariance (and again for any additional properties I want to obtain).

This isn't a huge hassle, but ideally I'd like to be able to obtain outputs as distribution objects instead of having the convert to tensor function run automatically, even if using model.predict with generators. Is there a way to do this?

My versions of TF and TFP are a bit out of date, I haven't updated them because my old code is currently working fine other than the one issue above. I will make a new venv later to test the issue on the latest version to see if it is already fixed.

My current versions are:
tf-nightly==2.3.0.dev20200611
tfp-nightly==0.11.0.dev20200624

Start by reproducing model.predict with a generator under the reported TensorFlow and TensorFlow Probability versions, then repeat with current versions to determine whether the conversion still occurs. Compare generator predictions with direct model calls; done means the supported prediction path preserves distribution objects or clearly documents the limitation.