gigl.src.common.modeling_task_specs.utils.infer_training_batch

gigl.src.common.modeling_task_specs.utils.infer.infer_training_batch(model: DistributedDataParallel | Module, training_batch: NodeAnchorBasedLinkPredictionBatch | RootedNodeNeighborhoodBatch | Data | HeteroData, gbml_config_pb_wrapper: GbmlConfigPbWrapper, device: device) → Dict[CondensedNodeType, Tensor]