Update convert_to_hf_gptneox.py

[leclem]

I propose this pull request for convert_to_hf_gptneox.py so it supports conversion to the HF format for the special case of non-distributed training, so with n_stages = 1