Problem about freeze the graph to pb file

thanks to the BertInferenceExample.java, i can do inference at java, but it’s google’s official bert, it use the MLP as the classifier, i added some layer as the classifier, and it works well.

When i try to use the freezeTrainedBert.py to freeze the graph, i choose the node name :loss/SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits
(i think it’s output should be a probability, so i choose this node, and i refered the code from “BertInferenceExample.java and freezeTrainedBert.py”, the author choose node “loss/Softmax”)

with freezeTrainedBert.py, i generated the pb file, but when i set the path to pb file, then run the “BertInferenceExample.java”, i get the follow Exception:

Blockquote
Exception in thread “main” java.lang.IllegalStateException: 1 Unprocessed nodes: [loss/SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits]
at org.nd4j.common.base.Preconditions.throwStateEx(Preconditions.java:641)
at org.nd4j.common.base.Preconditions.checkState(Preconditions.java:340)
at org.nd4j.imports.graphmapper.tf.TFGraphMapper.importGraph(TFGraphMapper.java:512)
at org.nd4j.imports.graphmapper.tf.TFGraphMapper.importGraph(TFGraphMapper.java:140)
at org.nd4j.imports.graphmapper.tf.TFGraphMapper.importGraph(TFGraphMapper.java:86)
at org.deeplearning4j.modelimportexamples.tf.advanced.bert.BertInferenceExample3.main(BertInferenceExample3.java:63)
Blockquote

I wonder how to solve this problem, thanks you so much.