20 lines
613 B
Python
20 lines
613 B
Python
|
"""Print tensor and segment metadata from an H3 direct-runner capture."""
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument("capture_dir")
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
for name in ("input.pt", "output.pt"):
|
||
|
|
payload = torch.load(f"{args.capture_dir}/{name}", map_location="cpu", weights_only=False)
|
||
|
|
print(name)
|
||
|
|
for key, value in payload.items():
|
||
|
|
if isinstance(value, torch.Tensor):
|
||
|
|
print(f" {key}: shape={tuple(value.shape)} dtype={value.dtype} finite={torch.isfinite(value).all().item()}")
|
||
|
|
else:
|
||
|
|
print(f" {key}: {value}")
|