glenn-jocher
commited on
Commit
•
77fb8ee
1
Parent(s):
22fb2b0
refactor dataloader
Browse files- utils/datasets.py +1 -1
utils/datasets.py
CHANGED
@@ -52,7 +52,7 @@ def create_dataloader(path, imgsz, batch_size, stride, opt, hyp=None, augment=Fa
|
|
52 |
pad=pad)
|
53 |
|
54 |
batch_size = min(batch_size, len(dataset))
|
55 |
-
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0,
|
56 |
dataloader = torch.utils.data.DataLoader(dataset,
|
57 |
batch_size=batch_size,
|
58 |
num_workers=nw,
|
|
|
52 |
pad=pad)
|
53 |
|
54 |
batch_size = min(batch_size, len(dataset))
|
55 |
+
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8]) # number of workers
|
56 |
dataloader = torch.utils.data.DataLoader(dataset,
|
57 |
batch_size=batch_size,
|
58 |
num_workers=nw,
|