Skip to content

Commit b0c83db

Browse files
authored
[DLMED] enhance doc-string of ToDevice (#2874)
Signed-off-by: Nic Ma <[email protected]>
1 parent e9f6e51 commit b0c83db

File tree

1 file changed

+6
-0
lines changed

1 file changed

+6
-0
lines changed

monai/transforms/utility/array.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1075,6 +1075,12 @@ class ToDevice(Transform):
10751075
Move PyTorch Tensor to the specified device.
10761076
It can help cache data into GPU and execute following logic on GPU directly.
10771077
1078+
Note:
1079+
If moving data to GPU device in the multi-processing workers of DataLoader, may got below CUDA error:
1080+
"RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing,
1081+
you must use the 'spawn' start method."
1082+
So usually suggest to set `num_workers=0` in the `DataLoader` or `ThreadDataLoader`.
1083+
10781084
"""
10791085

10801086
def __init__(self, device: Union[torch.device, str], **kwargs) -> None:

0 commit comments

Comments
 (0)