Fix a bug in inference code

[zheng-da]

Description of changes:
This PR reorganizes the code of performing inference of GNN and LM models.
Specifically, we split the nodes for inference based on locality. In this case, the embeddings are saved to local partitions via shared memory and we only need to run barrier before returning fron the inference function.
However, when computing LM embeddings, we split the nodes evenly to ensure all processes take roughly the same amount of time to compute LM embeddings. Otherwise, we will see a timeout in barrier in some processes. Because now the nodes are split evenly, we need to write data to remote memory. Before returning from the inference function, we need to call flush_data to ensure all data written to distributed memory can be read correctly.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.