File size: 355 Bytes
617d016 | 1 2 3 4 5 6 7 8 9 10 11 | """Minimal smoke test for HI-Mapper (no weights)."""
import torch
from hi_mapper import DivHiMapper
B, D = 2, 768
mapper = DivHiMapper(feat_dim=D, curvature=0.2, target_radius=1.0)
root, mid, leaves, hier_loss, prompt_loss, attr_loss = mapper(
torch.randn(B, 5, D), torch.randn(B, D)
)
print("root", tuple(root.shape), "hier_loss", float(hier_loss))
|