Thank you for responding. I am indeed using a very small model with only around 200k total parameters during training. The values above are when I set the model to eval mode and a batch size of 128. I also used another function torch.cuda.memory_summary() which printed the following
Validation batch  54  of  298
|===========================================================================|
|                  PyTorch CUDA memory summary, device ID 0                 |
|---------------------------------------------------------------------------|
|            CUDA OOMs: 0            |        cudaMalloc retries: 0         |
|===========================================================================|
|        Metric         | Cur Usage  | Peak Usage | Tot Alloc  | Tot Freed  |
|---------------------------------------------------------------------------|
| Allocated memory      |    4717 KB |    8050 KB |  159767 KB |  155050 KB |
|       from large pool |       0 KB |       0 KB |       0 KB |       0 KB |
|       from small pool |    4717 KB |    8050 KB |  159767 KB |  155050 KB |
|---------------------------------------------------------------------------|
| Active memory         |    4717 KB |    8050 KB |  159767 KB |  155050 KB |
|       from large pool |       0 KB |       0 KB |       0 KB |       0 KB |
|       from small pool |    4717 KB |    8050 KB |  159767 KB |  155050 KB |
|---------------------------------------------------------------------------|
| GPU reserved memory   |    8192 KB |    8192 KB |    8192 KB |       0 B  |
|       from large pool |       0 KB |       0 KB |       0 KB |       0 B  |
|       from small pool |    8192 KB |    8192 KB |    8192 KB |       0 B  |
|---------------------------------------------------------------------------|
| Non-releasable memory |    3474 KB |    3592 KB |  162857 KB |  159382 KB |
|       from large pool |       0 KB |       0 KB |       0 KB |       0 KB |
|       from small pool |    3474 KB |    3592 KB |  162857 KB |  159382 KB |
|---------------------------------------------------------------------------|
| Allocations           |     174    |     280    |    3635    |    3461    |
|       from large pool |       0    |       0    |       0    |       0    |
|       from small pool |     174    |     280    |    3635    |    3461    |
|---------------------------------------------------------------------------|
| Active allocs         |     174    |     280    |    3635    |    3461    |
|       from large pool |       0    |       0    |       0    |       0    |
|       from small pool |     174    |     280    |    3635    |    3461    |
|---------------------------------------------------------------------------|
| GPU reserved segments |       4    |       4    |       4    |       0    |
|       from large pool |       0    |       0    |       0    |       0    |
|       from small pool |       4    |       4    |       4    |       0    |
|---------------------------------------------------------------------------|
| Non-releasable allocs |       5    |       8    |    1273    |    1268    |
|       from large pool |       0    |       0    |       0    |       0    |
|       from small pool |       5    |       8    |    1273    |    1268    |
|===========================================================================|

which also shows that my model is using 4717 KB (allocated memory) at test time which is equivalent to 0.004499GB.
I have two last questions please.


A batch size of 128 prints torch.cuda.memory_allocated: 0.004499GB whereas increasing it to 1024 prints torch.cuda.memory_allocated: 0.005283GB. Can I confirm that the difference of approximately 1MB is only due to the increased batch size?


And is there a reason why nvidia-smi goes from 1349MiB to 1355MiB when going from a batch size of 128 to 1024? The increase is not consistent with the 1MB I got from the previous point. Does the memory used by CUDA context also vary based on the batch size?