Error when converting data to heterogenous

[EvanUp]

I have a large custom dataset and I want to add a single edge attribute to my model and then feed the heterogenous data to the model in batches using neighborloader. I tried using the .to_heterogenous and .to_hetero methods to convert my dataset to a heterogenous dataset, but I'm getting the following error when I try to pass batches to the model:

Traceback (most recent call last):
  File "/usr0/home/emwillia/.pyenv/versions/miniconda3-4.7.12/lib/python3.7/site-packages/torch/fx/graph_module.py", line 505, in wrapped_call
    return cls_call(self, *args, **kwargs)
  File "/usr0/home/emwillia/.pyenv/versions/miniconda3-4.7.12/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "<eval_with_key_1>", line 3, in forward
    conv2__0 = torch.add(None, None)
TypeError: add(): argument 'input' (position 1) must be Tensor, not NoneType

Call using an FX-traced Module, line 3 of the traced Module's generated forward function:
def forward(self, x, edge_index):
    conv2__0 = torch.add(None, None)

~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    return {'0': conv2__0}

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr0/home/emwillia/.pyenv/versions/miniconda3-4.7.12/lib/python3.7/site-packages/torch/fx/graph_module.py", line 513, in wrapped_call
    raise e.with_traceback(None)
TypeError: add(): argument 'input' (position 1) must be Tensor, not NoneType

Any help would be greatly appreciated! :) I've included a reprex below.

Karate Club Reprex:

## torch.__version__ == '1.9.1+cu111'
## torch_geometric.__version__ == '2.0.2'

import torch_geometric.transforms as T
import networkx as nx
from torch_geometric.utils.convert import from_networkx
from torch_geometric.loader import NeighborLoader
from torch_geometric.nn import SAGEConv, to_hetero
import torch

# make a custom dataset and convert to heterogenous
G = nx.karate_club_graph()
dataset = from_networkx(G)
dataset.x = torch.FloatTensor(len(G.nodes()), 5).uniform_(1, 5)
dataset.train_mask = torch.randint(0, 2, (34,))
recode_labels = {'Mr. Hi': 0, 'Officer':1}
dataset.y = torch.tensor([recode_labels[i] for i in dataset.club]).long()
dataset = dataset.to_heterogeneous(edge_type = torch.randint(0,2, (dataset.num_edges,)))
dataset = T.ToUndirected()(dataset)
dataset = T.AddSelfLoops()(dataset)
dataset = T.NormalizeFeatures()(dataset)

print(dataset)

karate_loader = NeighborLoader(dataset, input_nodes = ('0', torch.ones(len(dataset['0'].train_mask)).long()),
                        num_neighbors=[5, 5], batch_size=4, shuffle=False,
                        num_workers=0)

class GNN(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = SAGEConv((-1, -1), hidden_channels)
        self.conv2 = SAGEConv((-1, -1), out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        return x
    
model = GNN(4, 2)
model = to_hetero(model, dataset.metadata(), aggr='sum').to('cpu')
batch = next(iter(karate_loader))
batch = batch.to('cpu')
out = model(batch.x_dict, batch.edge_index_dict)

[EvanUp]

Oh, figured it out... had two issues:

1. node names and edge names need to be different
2. my mask wasn't a boolean