Conditional Tensor Shape for PyTorch Linear Layers

I have a dataset loaded into a dataloader with the following features.

x: (node_features * batch_size, 4)
edge_index: (2, edge_size)
y: (node_features * batch_size, )
z: (node_features * batch_size, )
boolean: (batch_size, )

I am using a Graph Neural Network to learn an output tensor p of shape (node_features * batch_size). I have access to all data during training, but not during evaluation, validation or inference time. In that time I only have access to x and edge_index.

How do I learn from all the data in the dataloader during training, but only use x and edge_index for testing?

One might think, simply concatenating the other values to the output before passing it to the linear layer would work, but because the dimensions change during inference time, the linear layer will complain about the un-concatenated data. How do I solve this issue?