print(f'train_mask = {data.train_mask.shape}')
print(data.train_mask)