Function sanitize_pair

flowmap/visualization/drawing/types.py:61–67

(pair: Pair, device: torch.device)

def sanitize_pair(pair: Pair, device: torch.device) -> Float[Tensor, "2"]:
    if isinstance(pair, Tensor):
        pair = pair.type(torch.float32).to(device)
    else:
        pair = torch.tensor(pair, dtype=torch.float32, device=device)
    assert pair.shape == (2,)
    return pair