Change PackSegments to ensure consistent behavior between CPU and GPU
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/31673
Reviewed By: Wakeupbuddy, BIT-silence
Differential Revision: D18925762
fbshipit-source-id: e0c318e97f69b14a54f43c176af57d98fbc16c9f