Hi,
How do you measure the gpu memory usage? You should be using torch.cuda.memory_allocated to get the memory actually used by Tensors. Note that this number only counts Tensors, and not the memory required by the cuda driver at initialization.
If you check via nvidia-smi, you will see a larger number because we have a special allocator for speed reasons that do not return the memory back to the driver when Tensors are freed.