Compute the GPU memory usage for the current device (GB).
()
| 51 | |
| 52 | |
| 53 | def gpu_mem_usage(): |
| 54 | """ |
| 55 | Compute the GPU memory usage for the current device (GB). |
| 56 | """ |
| 57 | if torch.cuda.is_available(): |
| 58 | mem_usage_bytes = torch.cuda.max_memory_allocated() |
| 59 | else: |
| 60 | mem_usage_bytes = 0 |
| 61 | return mem_usage_bytes / 1024 ** 3 |
| 62 | |
| 63 | |
| 64 | def cpu_mem_usage(): |