This seems as you describe.
Memory is returned to the torch cache once the tensors using it go out of scope.
It depends on your optimizer, but things like tracking statistics for adaptivity (eg in Adam) or momentum use additional memory. These are allocated once, so it should not increase over time.

In general, I would advise to no use cuda_malloc directly but instead create a new torch tensor. As long as you keep a reference to it, the memory will be reserved for you and can be accessed via data_ptr() (for a void*) or data_ptr<T>() for a T*.
Best regards
Thomas