Thanks!
As you can see in the memory_summary(), PyTorch reserves ~2GB so given the model size + CUDA context + the PyTorch cache, the memory usage is expected:
| GPU reserved memory   |    2038 MB |    2038 MB |    2038 MB |       0 B  |
|       from large pool |    2036 MB |    2036 MB |    2036 MB |       0 B  |
|       from small pool |       2 MB |       2 MB |       2 MB |       0 B  |

If you want to release the cache, use torch.cuda.empty_cache(). This will synchronize your code thus slowing it down, but would allow other applications to use this memory.