Skip to content

Commit eb5e3f7

Browse files
committed
fix: device mismatch in stats xp.take index
1 parent e437691 commit eb5e3f7

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

src/earthkit/meteo/stats/array/quantiles.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -77,9 +77,9 @@ def iter_quantiles(
7777
f = (m - 1) * q
7878
j = int(f)
7979
x = f - j
80-
quantile = xp.take(arr, xp.asarray(j), axis=axis)
80+
quantile = xp.take(arr, xp.asarray(j, device=device), axis=axis)
8181
quantile *= 1 - x
82-
tmp = xp.take(arr, xp.asarray(min(j + 1, m - 1)), axis=axis)
82+
tmp = xp.take(arr, xp.asarray(min(j + 1, m - 1), device=device), axis=axis)
8383
tmp *= x
8484
quantile += tmp
8585
quantile[xp.reshape(missing, quantile.shape)] = xp.nan

0 commit comments

Comments
 (0)