(device=None)
| 69 | |
| 70 | |
| 71 | def get_cuda_free_memory_gb(device=None): |
| 72 | if device is None: |
| 73 | device = gpu |
| 74 | |
| 75 | memory_stats = torch.cuda.memory_stats(device) |
| 76 | bytes_active = memory_stats['active_bytes.all.current'] |
| 77 | bytes_reserved = memory_stats['reserved_bytes.all.current'] |
| 78 | bytes_free_cuda, _ = torch.cuda.mem_get_info(device) |
| 79 | bytes_inactive_reserved = bytes_reserved - bytes_active |
| 80 | bytes_total_available = bytes_free_cuda + bytes_inactive_reserved |
| 81 | return bytes_total_available / (1024 ** 3) |
| 82 | |
| 83 | |
| 84 | def move_model_to_device_with_memory_preservation(model, target_device, preserved_memory_gb=0): |
no outgoing calls
no test coverage detected