import torch, sys, os
os.chdir(os.path.dirname(os.path.abspath(__file__)))
pth = 'best_v13_mild_r2.pth'
ckpt = torch.load(pth, map_location='cpu', weights_only=False)
print(f"Type: {type(ckpt).__name__}")
if isinstance(ckpt, dict):
    if 'model' in ckpt:
        print("Has 'model' key")
        ckpt = ckpt['model']
    if 'state_dict' in ckpt:
        print("Has 'state_dict' key")
        ckpt = ckpt['state_dict']
    print(f"\nTotal keys: {len(ckpt)}")
    for k, v in sorted(ckpt.items()):
        shape = tuple(v.shape) if hasattr(v, 'shape') else 'scalar'
        print(f"  {k:50s} {str(shape)}")
else:
    print("Not a dict, top-level keys unknown")