fix mypy

pytorch · Oct 21, 2024 · cdbd920 · cdbd920
1 parent b7b18a2
commit cdbd920
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/torchdata/nodes/pin_memory.py b/torchdata/nodes/pin_memory.py
@@ -3,7 +3,7 @@
 import queue
 import threading
 
-from typing import Iterator
+from typing import Iterator, Optional, Union
 
 import torch
 import torch.multiprocessing
@@ -16,12 +16,12 @@
 
 
 def _pin_memory_loop(
-    source: queue.Queue,
+    source: BaseNode,
     q: queue.Queue,
     semaphore: threading.BoundedSemaphore,
     stop_event: threading.Event,
-    device_id: str,
-    device: str,
+    device_id: Union[int, str],
+    device: Optional[str],
 ):
     # this is fork of from torch.utils.data._utils.pin_memory import _pin_memory_loop
     # to remove the index tuples