From ec8b5f768dc566ecf36f54fede940860aecc8268 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Val=C3=A9rian=20Rey?= <valerian.rey@gmail.com>
Date: Wed, 4 Feb 2026 04:12:00 +0100
Subject: [PATCH] test: Fix fork_rng context

It should now work with other devices than cpu and cuda
---
 tests/utils/contexts.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/utils/contexts.py b/tests/utils/contexts.py
index dc508130..db2d54eb 100644
--- a/tests/utils/contexts.py
+++ b/tests/utils/contexts.py
@@ -10,7 +10,7 @@
 
 @contextmanager
 def fork_rng(seed: int = 0) -> Generator[Any, None, None]:
-    devices = [DEVICE] if DEVICE.type == "cuda" else []
+    devices = [] if DEVICE.type == "cpu" else [DEVICE]
     with torch.random.fork_rng(devices=devices, device_type=DEVICE.type) as ctx:
         torch.manual_seed(seed)
         yield ctx