Add DeviceType for NumpyToTensor

Signed-off-by: Daniel Sperber <[email protected]>

Author: Daraan

rllib/connectors/common/numpy_to_tensor.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, TYPE_CHECKING
 
 import gymnasium as gym
 
@@ -12,6 +12,9 @@
 from ray.rllib.utils.typing import EpisodeType
 from ray.util.annotations import PublicAPI
 
+if TYPE_CHECKING:
+    from ray.rllib.utils.typing import DeviceType
+
 
 @PublicAPI(stability="alpha")
 class NumpyToTensor(ConnectorV2):
@@ -59,7 +62,7 @@ def __init__(
         input_action_space: Optional[gym.Space] = None,
         *,
         pin_memory: bool = False,
-        device: Optional[str] = None,
+        device: Optional["DeviceType"] = None,
         **kwargs,
     ):
         """Initializes a NumpyToTensor instance.