Compute the GPU memory usage for the current device (MB).
()
| 41 | |
| 42 | |
| 43 | def gpu_mem_usage(): |
| 44 | """ |
| 45 | Compute the GPU memory usage for the current device (MB). |
| 46 | """ |
| 47 | mem_usage_bytes = torch.cuda.max_memory_allocated() |
| 48 | return mem_usage_bytes / (1024 * 1024) |
| 49 | |
| 50 | |
| 51 | class AverageMeter: |