Including global graph information into the loss function

Dear ignnition team,

In our problem we are training a GNN for a regression task - a subset of nodes is labeled by a float value and those labels are learned using a neural network as a readout model. The whole GNN model is trained based on the MSE loss between the labels and the predictions for the mentioned subset of nodes and works well!

We would like to create a new loss function that incorporates some physical laws related to our problem. In each training step, after the predictions are generated for all of the labeled nodes, we would like to add an additional term to the MSE between the labels and the predictions. That additional term would multiply all of the predictions with some coefficients (different for every node), and sum all of the obtained values. So the goal would be to minimize that sum along with the MSE.

Is something like this possible to implement in the ignnition framework? I'm not even sure is it consistent with the logic for creating the mini-batches - I guess that the requirement here would be to have all of the nodes from one training sample in the same mini-batch.

Thanks!

Ognjen