gigl.src.common.modeling_task_specs.utils.infer_root_embeddings

gigl.src.common.modeling_task_specs.utils.infer.infer_root_embeddings(model: DistributedDataParallel | Module, graph: Data | HeteroData, root_node_indices: LongTensor, gbml_config_pb_wrapper: GbmlConfigPbWrapper, device: device) → Tensor