Print warnings/errors for large swap space (#123)

2023-05-23 18:22:26 -07:00
parent a283ec2eec
commit aedba6d5ec
3 changed files with 34 additions and 0 deletions
--- a/cacheflow/utils.py
+++ b/cacheflow/utils.py
@@ -24,8 +24,10 @@ class Counter:


 def get_gpu_memory(gpu: int = 0) -> int:
+    """Returns the total memory of the GPU in bytes."""
    return torch.cuda.get_device_properties(gpu).total_memory


 def get_cpu_memory() -> int:
+    """Returns the total CPU memory of the node in bytes."""
    return psutil.virtual_memory().total