Get gpu memory used by current JAX script

get_jax_gpu_mem.py

import subprocess as subp
import re
if jax.default_backend()=='gpu':
    # Step 2: Call nvidia-smi and get the output
    result = subp.run([
        'nvidia-smi', 
        '--query-compute-apps=pid,used_memory', 
        '--format=csv,noheader,nounits'
    ], stdout=subp.PIPE, text=True)
    
    # Step 3: Get the output
    output = result.stdout.strip()
    
    # Step 4: Process the output
    memory_usage = {}
    for line in output.split('\n'):
        pid, mem = line.split(',')
        memory_usage[int(pid.strip())] = int(mem.strip())
    
    # Step 5: Print the memory usage for current PID
    print(f"GPU memory: {memory_usage[os.getpid()]} MiB")