| #!/usr/bin/env python3 | |
| import json, torch | |
| o={"ok":bool(torch.cuda.is_available())} | |
| if torch.cuda.is_available(): | |
| f,t=torch.cuda.mem_get_info() | |
| G=1024**3 | |
| o.update({ | |
| "free_gib":f/G, | |
| "total_gib":t/G, | |
| "used_gib":(t-f)/G, | |
| "free_ratio":f/t, | |
| "device_name":torch.cuda.get_device_name(0), | |
| "torch_allocated_gib":torch.cuda.memory_allocated()/G, | |
| "torch_reserved_gib":torch.cuda.memory_reserved()/G, | |
| }) | |
| print(json.dumps(o,indent=2)) | |
| node-7.example.invalid SystemExit(0 if o["ok"] else 2) | |