Just create a flag to check whether the tensor has been allocated or not. Allocate it once.
Another option is just to delete it when you stop training.
Anyway I think it’s not bad it stays as a buffer. It’s somehow a necessary parameter to run inference. What would you do for exporting the model, keep telling everyone those values? Seems a bit annoying to me.