Thank you @polvalente for mentioning garbage collection.
If I set the :garbage_collect option to true with Axon.Loop.run then the memory usage is low and it doesn’t increase. If I don’t set it, then it grows continuously during training.
This solves the issue.
