[Layers] Fixed the seg fault error when running with more than 4 ranks #424

Closed abenmao closed 1 month ago