kalpeshjp89:

In test method, it iteratively consumes more and more memory until all 16GB and then throwing error.


This points towards a potential storage of a tensor, which is attached to the computation graph.
I cannot find anything obviously wrong in the posted code snippet, but could you check, if you are storing some outputs with a valid .grad_fn in a list/dict etc.?