annealed mean implementation

[laaRaa]

Hi,

I have a question regarding the annealed mean implementation (PyTorch code). If I understand correctly this step is implemented, for testing, as an additional convolutional layer applied to the distribution tensor (after the softmax layer). If this is correct, I do not understand how can the function f_T (equation (5) in the paper) be implemented with a convolutional layer. Or is it only for the case where the temperature T is equal to 1 (hence taking the mean and not the annealed mean)?

Thanks in advance!