# Forward pass
    h, z = model(data.x, data.edge_index)