Can use graph cut to multi-label segment?

[tanjia123456]

Hello, thanks of your contribution about such good work.
I am using graph neural network to parcel brain region, however the result is not good, I want to use graph cut as post process.
The inputs of my current model are: adjacency matrix (10242, 10242), feature matrix (10242, 6), label matrix (10242). The output is the probability y that each node belongs to a label, and its dimension is (10242, 36)
I want to use graph cut to update y for better performance.

I have a few questions about your code:
First, it seems that your code is only for two categories. Can you do multi label tasks?
Second, if I want to do post-processing, what should my input be?