[Bugfix] Fix a small device bug

[szhengac]

Initialize all-gather tensor on gpu.

Checklist

• PR's title starts with a category (e.g. [Bugfix], [Model], [Tutorial], etc)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented

[chhzh123]

This line also uses .to(device). Maybe we can also change it?
https://github.com/awslabs/slapo/blob/main/slapo/framework_dialect/deepspeed/pipeline.py#L320

[szhengac]

This line also uses .to(device). Maybe we can also change it? https://github.com/awslabs/slapo/blob/main/slapo/framework_dialect/deepspeed/pipeline.py#L320

Nice catch!

[comaniac]

Thanks @szhengac @chhzh123