Fixed device mismatch between variable

[sonhm3029]

📝 Description

• Provide a clear summary of the changes and the issue that has been addressed.
• 🛠️ Fixes # (issue number)

✨ Changes

Select what type of change your PR is:

• 🚀 New feature (non-breaking change which adds functionality)
• 🐞 Bug fix (non-breaking change which fixes an issue)
• 🔄 Refactor (non-breaking change which refactors the code base)
• ⚡ Performance improvements
• 🎨 Style changes (code style/formatting)
• 🧪 Tests (adding/modifying tests)
• 📚 Documentation update
• 📦 Build system changes
• 🚧 CI/CD configuration
• 🔧 Chore (general maintenance)
• 🔒 Security update
• 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)

✅ Checklist

Before you submit your pull request, please make sure you have completed the following steps:

• 📚 I have made the necessary updates to the documentation (if applicable).
• 🧪 I have written tests that support my changes and prove that my fix is effective or my feature works (if applicable).
• 🏷️ My PR title follows conventional commit format.

For more information about code review checklists, see the Code Review Checklist.

[rajeshgangireddy]

References
Discussion : #2749
Issue : #2923

[rajeshgangireddy]

Will still work for devices that are neither cuda or cpu. For example, when we have intel GPUs, the device is called xpu. It would be nice to solve this such that it works independent of the device used.