Thanks for answering!
I just tried that, however, I’m getting problems with the tensor broadcasting when calculating eta:
These are the tensor shapes for the model:
And for the guide:
