Memory usage increasing with each epoch

setting environment variables like this

ENV["JULIA_CUDA_HARD_MEMORY_LIMIT"] = "3.0GiB"

or

ENV["JULIA_CUDA_SOFT_MEMORY_LIMIT"] = "3.0GiB"

can mitigate the issue