Memory leak for TFP distribution

[fotisdr]

Hi,

I'm facing a memory leak issue with TFP distributions. Specifically, I have a custom function that samples from a Categorical distribution and returns the output:

def compute_categorical_distribution(classes_pred,dtype=np.int8,validate_args=True):
  # Define the distribution
  distr = tfp.distributions.Categorical(probs=classes_pred, validate_args=validate_args)
  # Sample from the distribution
  output = distr.sample().numpy().astype(dtype)
del distr
K.clear_session()
gc.collect()
return output

I was expecting that, when I call this function, I would get the output array back and the memory allocation of the distribution would be freed up. However, this does not seem to be the case, as for every consecutive call of this function the memory usage keeps increasing, indicating that a new tfp distribution is created each time and remains in the memory. Is there a way to clear up the memory allocation of the distribution? As you can see above, I tried del distr, tensorflow.keras.backend.clear_session() and gc.collect() but none of these seems to make any difference. I am using tfp version 0.19.0.

Thanks!
Fotis