This error seems to be raised by Python’s import system and apparently Python isn’t able to find the definition of NusDatasetGCN. Are you importing it from another file or define everything in a single script?