Typeerror: batch must contain tensors, numbers, dicts or lists

Hello together,

I’m trying to train a GNN using the MNIST dataset which I previously turn into a slic segmentation using the Compose class:

import torch 
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.transforms as T
from torch_geometric.transforms import ToSLIC

trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
trans = T.Compose([T.ToTensor(), ToSLIC(n_segments=75)])

train_set = dset.MNIST(root='/data', train=True, transform=trans, download=True)
test_set = dset.MNIST(root='/data', train=False, transform=trans, download=True)

batch_size = 100

Afterwards I try to train my model however I have issues that I get receive the error that I’m passing an object and not a tensor. My model training looks somewhat like this:

from torch.nn import Linear
import torch.nn.functional as F
from torch_geometric.nn import GCNConv, TopKPooling, global_mean_pool
from torch_geometric.nn import global_mean_pool as gap, global_max_pool as gmp

embedding_size = 64

class GCN(torch.nn.Module):
    def __init__(self):
        super(GCN, self).__init__()
        self.init_conv = GCNConv(num_features, 64)
        self.conv1 = GCNConv(num_features, 64)
        self.conv2 = GCNConv(num_features, 64)
        self.conv3 = GCNConv(num_features, 64)
        self.out = Linear(embedding_size*2, 10)
    def forward(self, x, edge_index, batch_index):
        hidden = self.initial_conv(x, edge_index)
        hidden = F.tanh(hidden)

        hidden = self.conv1(hidden, edge_index)
        hidden = F.tanh(hidden)
        hidden = self.conv2(hidden, edge_index)
        hidden = F.tanh(hidden)
        hidden = self.conv3(hidden, edge_index)
        hidden = F.tanh(hidden)
        hidden = torch.cat([gmp(hidden, batch_index),
                           gap(hidden, batch_index)])
        out = self.out(hidden)
        return out, hidden      

model = GCN()

loss = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(),lr=0.0001)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)

data_size = len(data)

train_loader = torch.utils.data.DataLoader(data[:int(data_size*0.8)],
                   batch_size = NUM_PER_BATCH, shuffle=True)

test_loader = torch.utils.data.DataLoader(data[int(data_size*0.8):],
                   batch_size = NUM_PER_BATCH, shuffle=True)

def train(train_loader):
    for batch in train_loader:
        pred, embedding = model(batch.x.float(), batch.pos.float(), batch.batch)
        loss = torch.sqrt(loss(pred, batch.y))
    return loss, embedding

loss = []
for epochs in range(10):
    loss, h = train(train_loader)
    if epoch % 2 == 0:
        print("Loss: {}".format(loss))

What I also noticed is that ToSlic does not store the edge_index which I assume has to be precalculated (and might cause the error). For guidance I would be very grateful.