fix: Bugfix in shape analysis for multi-GPU systems

[gs-olive]

Description

• Shape analysis code in partitioning defaults dry-run tensors to cuda:0 despite user-specified devices
• This leads to errors about device casting for internal tensors, which users cannot cast
• Add GPU-ID function arguments in functions to generate new tensors on the user-specified (or default) device

Fixes #1764

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• [ x ] My code follows the style guidelines of this project (You can use the linters)
• [ x ] I have performed a self-review of my own code
• [ x ] I have commented my code, particularly in hard-to-understand areas and hacks
• [ x ] I have made corresponding changes to the documentation
• [ - ] I have added tests to verify my fix or my feature
  • Verified locally [requires multi-GPU to test]
• [ x ] New and existing unit tests pass locally with my changes
• [ x ] I have added the relevant labels to my PR in so that relevant reviewers are notified

[peri044]

LGTM