Get current GPU VRAM usage in MB.
()
| 509 | # --------------------------------------------------------------------------- |
| 510 | |
| 511 | def get_vram_usage() -> Dict[str, float]: |
| 512 | """Get current GPU VRAM usage in MB.""" |
| 513 | import torch |
| 514 | if not torch.cuda.is_available(): |
| 515 | return {"allocated_mb": 0, "reserved_mb": 0, "peak_mb": 0} |
| 516 | return { |
| 517 | "allocated_mb": torch.cuda.memory_allocated() / 1e6, |
| 518 | "reserved_mb": torch.cuda.memory_reserved() / 1e6, |
| 519 | "peak_mb": torch.cuda.max_memory_allocated() / 1e6, |
| 520 | } |
| 521 | |
| 522 | |
| 523 | # --------------------------------------------------------------------------- |