test: Fix fork_rng context (#556)

* It should now work with other devices than cpu and cuda.

Author: ValerianRey

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