p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] utils/train_utils.py (python)
def print_mem_stats(label: str):
  max_logging.log(f"
Memstats: {label}:")
  try:
    for d in jax.local_devices():
      stats = d.memory_stats()
      used = round(stats["bytes_in_use"] / 2**30, 2)
      limit = round(stats["bytes_limit"] / 2**30, 2)
      max_logging.log(f"	Using (GB) {used} / {limit} ({used/limit:%}) on {d}")
  except (RuntimeError, KeyError, TypeError) as ex:
    max_logging.log(f"	Memstats unavailable, error: {ex}")
[content] train_tokenizer.py (python)
def print_mem_stats(label: str):
  max_logging.log(f"
Memstats: {label}:")
  try:
    for d in jax.local_devices():
      stats = d.memory_stats()
      used = round(stats["bytes_in_use"] / 2**30, 2)
      limit = round(stats["bytes_limit"] / 2**30, 2)
      max_logging.log(f"	Using (GB) {used} / {limit} ({used/limit:%}) on {d}")
  except (RuntimeError, KeyError, TypeError) as ex:
    max_logging.log(f"	Memstats unavailable, error: {ex}")