[Pytorch Geometric] Question over how to pass in data from DataLoader

Hello, I’m a newbie to playing around with pytorch geometric, and I’ve made the following toy code (it’s based on a real problem I’m working on):

import pandas as pd
import numpy as np
import argparse
import torch
from torch_geometric.nn import EdgeConv
from torch_geometric.data import Data, DataLoader
from torch.nn import Sequential, Linear

X = torch.randn(1000,4,4)
Y = torch.randn(1000,5000)

graphs = []
for vector in X:
    x = torch.tensor([[i] for i in vector.flatten()])
    edges = []
    node_index = 0
    for i in range(4):
        for j in range(4):
            if (j != 3): edges.append( [node_index, node_index + 1] )
            if (j != 0): edges.append( [node_index, node_index - 1] )
            if (i != 3): edges.append( [node_index, node_index + 4] )
            if (i != 0): edges.append( [node_index, node_index - 4] )
            node_index += 1
    edges = torch.tensor(edges, dtype = torch.long).t().contiguous()
    graphs.append( Data(x = x, edge_index = edges) )
loader = DataLoader(graphs, batch_size = 256)

m1 = Sequential(
    Linear(2,1)
)

m2 = Sequential(
    Linear(2,1)
)

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = EdgeConv(m1, aggr = 'mean')
        self.conv2 = EdgeConv(m2, aggr = 'mean')
    
    def forward(self, data):
        x, edge_index = data.x, data.edge_index

        x = self.conv1(x, edge_index)
        x = self.conv2(x, edge_index)
        x = x.t()

        return Linear(16, 4624)(x)

model = Net()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
loss = torch.nn.MSELoss()

model.train()
for epoch in range(200):
    optimizer.zero_grad()
    i = 0
    for train_x in loader:
        optimizer.zero_grad()
        out = model(train_x)
        i += 1
    print(i)

When I run this code using torch == 1.4.0 and torch_geometric == 1.6.0, with Python 3.7, I get the following error:

RuntimeError: size mismatch, m1: [1 x 4096], m2: [16 x 4624]

Could someone educate me as to what I’m doing wrong? I know that if I run the Net() on an individual graph it works just fine, however it appears the issue is with how PyTorch Geometric is passing data from the batched graphs. I know it does something different when creating minibatches using DataLoader to allow for efficient mini-batch computation, but it’s not clear from the documentation or examples how exactly we’re supposed to work with these batches. I’ve looked high and low to try and find samples showing how to work with PyTorch Geometric’s DataLoader class, but I can’t seem to find any good references.

Apologies if this is a very elementary question, I’m fairly new at PyTorch in general (I’m coming from TensorFlow) and while I see how some of the concepts line up I don’t think I know enough to effectively debug this problem on my own. Any help would be greatly appreciated.

Hi, I’m having exactly the same problem here. Did you find a solution to that? Thanks