Free memory in universal checkpointing tests (#6693)
Tests in universal checkpointing were not freeing the engine after use
when `reuse_dist_env` was set to `True`, leading to memory leaks.
This PR ensure freeing the engine in the tests and enables
`reuse_dist_env`.
Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>