I need to calculate the exponential moving average on a tensor inside of a custom loss function but searching on the internet it looks like tensorflow only supports calculating the moving average on a model’s parameters.

My model uses LSTMs and both the input and the output are 3D arrays (the ouput has the exact same shape of the input).
Inside my custom loss I calculate the mse getting a 2D tensor (the third dimension “collapses”), now I need to calculate the moving average along the axis 1, but I cannot find any way to do this, I tried converting the tensor to a ndarray to do it “manually” but I can’t find the right way to that either (I can’t run the code in eager mode for performance issues, so .numpy() doesn’t work and I tried using a tf.Session() but I get the error “AttributeError: module ‘tensorflow’ has no attribute ‘Session'”).

Am I getting the way tf.train.ExponentialMovingAverage works wrong?
Could I use it to calculate the moving average on a tensor?
If not, what could be a way to avoid the tf.Session() error to calculate the ema manually?
At last, does anybody know if there is any built in way to calculate the ema on a tensor in tensorflow?