#!/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)