def METHOD_NAME(self, data: Union[Data, Dict[str, torch.Tensor]]) -> torch.Tensor:
    if "batch" in data:
        batch = data["batch"]
    else:
        batch = data["pos"].new_zeros(data["pos"].shape[0], dtype=torch.long)
    # Create graph
    edge_index = radius_graph(data["pos"], self.max_radius, batch, max_num_neighbors=len(data["pos"]) - 1)
    edge_src = edge_index[0]
    edge_dst = edge_index[1]
    # Edge attributes
    edge_vec = data["pos"][edge_src] - data["pos"][edge_dst]
    return batch, data["x"], edge_src, edge_dst, edge_vec