Skip to content
This repository was archived by the owner on Sep 10, 2025. It is now read-only.

Commit 85be3b7

Browse files
authored
use 'cpu' instead of None for Iterator (#745)
1 parent 8b58a22 commit 85be3b7

File tree

1 file changed

+7
-0
lines changed

1 file changed

+7
-0
lines changed

torchtext/data/iterator.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
import logging
55

6+
import torch
67
from .utils import RandomShuffler
78
from .batch import Batch
89
from .dataset import Dataset
@@ -65,6 +66,12 @@ def __init__(self, dataset, batch_size, sort_key=None, device=None,
6566
+ " or passing a string as an argument. This behavior will be"
6667
+ " deprecated soon and currently defaults to cpu.")
6768
device = None
69+
70+
if device is None:
71+
device = torch.device('cpu')
72+
elif isinstance(device, str):
73+
device = torch.device(device)
74+
6875
self.device = device
6976
self.random_shuffler = RandomShuffler()
7077

0 commit comments

Comments
 (0)