Precision issue on mobile devices

[tianhaoz95]

Describe the bug

I implemented the loading of safetensor and then inference for qwen 2.5 0.5B model, and tested with request "what is 1+1? only answer with numbers" it correctly output "2" with both burn::backend::NdArray and burn::backend::Wgpu as backend on my macbook, but on Android device, CPU works fine, but vulkan outputs garbage, on iOS, CPU works fine, but metal outputs garbage.

for mobile device i did extra setup:

#[cfg(any(target_os = "ios", target_os = "macos"))]
pub type GraphicsApi = burn_wgpu::graphics::Metal;
#[cfg(target_os = "android")]
pub type GraphicsApi = burn_wgpu::graphics::Vulkan;
#[cfg(not(any(target_os = "ios", target_os = "macos", target_os = "android")))]
pub type GraphicsApi = burn_wgpu::graphics::AutoGraphicsApi;

...

#[cfg(feature = "high_perf")]
{
    if GPU_SETUP.get().is_none() {
        let _setup = burn_wgpu::init_setup_async::<backend_setup::GraphicsApi>(
            &device,
            burn_wgpu::RuntimeOptions::default(),
        )
        .await;
        let _ = GPU_SETUP.set(());
    }
}

To Reproduce

1. given model implementation https://gist.github.com/tianhaoz95/f8e2fd444d8b2954e4f57f60ca2d3812
2. run it on desktop with prompt "what is 1+1? only answer with numbers" , outputs "2" for CPU and GPU
3. run it on Android, CPU outputs "2" and vulkan output garbage

Expected behavior

They should output same value or similar content

Screenshots

NA

Desktop (please complete the following information):

• OS: Mac
• Browser Chrome
• Version 26

Smartphone (please complete the following information):

• Device: Android
• OS: 16
• Browser Chrome
• Version NA

Additional context