Source code for cogdl.wrappers.data_wrapper.heterogeneous.heterogeneous_gnn_dw

from .. import DataWrapper


[docs]class HeterogeneousGNNDataWrapper(DataWrapper):
    def __init__(self, dataset):
        super(HeterogeneousGNNDataWrapper, self).__init__(dataset=dataset)

        self.dataset = dataset

[docs]    def train_wrapper(self):
        return self.dataset

[docs]    def val_wrapper(self):
        return self.dataset

[docs]    def test_wrapper(self):
        return self.dataset