This PR adds a script argument --post_pool_filter that allows us to specify a list of ranks whose cells should be used to make the final prediction. By default, cells of all ranks are used. By setting post_pool_filter 0, we can now ensure that only atom representations would be used.
The main purpose of this argument is to be able to simulate EGNN with our codebase. We want to add the supercell to get fully-connected ranks but do not want the supercell itself to influence the final prediction.
This PR adds a script argument
--post_pool_filterthat allows us to specify a list of ranks whose cells should be used to make the final prediction. By default, cells of all ranks are used. By settingpost_pool_filter 0, we can now ensure that only atom representations would be used.The main purpose of this argument is to be able to simulate EGNN with our codebase. We want to add the supercell to get fully-connected ranks but do not want the supercell itself to influence the final prediction.