Fixed iterator to only store data for that rank.

[hariharan-devarajan]

• we fix reconfiguration logic to build rank map for train and valid.
• Now we need reconfigure per epoch to shuffle the list.