Fill most MPS memory with blocks by default

9 months ago · 4534c6ed1f
parent b79c357d5d
commit 4534c6ed1f
1 changed files with 1 additions and 1 deletions
--- a/src/petals/server/server.py
+++ b/src/petals/server/server.py
@ -288,7 +288,7 @@ class Server:
        elif self.device.type == "cuda":
            total_memory = torch.cuda.get_device_properties(self.device).total_memory
        else:
-            total_memory = psutil.virtual_memory().available
+            total_memory = psutil.virtual_memory().total

        gib = 1024**3
        # Estimate of GPU memory used in rpc_backward (2 GiB for BLOOM, proportional for other models)