Some of the tests were flaky, and I suspect the problem came from unreliable initialization in burn-wgpu. The problem comes from the fact that zeros was achieved by multiplying an empty tensor by zero. However, the empty tensor might have invalid numbers in it, or NaN, which would make the zeros or ones function return NaN as well. The fix is pretty simple: create a Full kernel and use it to initialize tensors. This reduces the number of kernels launched and avoids uninitialized memory from causing numerical problems. I wasn't able to reproduce the flaky test after the change.