Skip to content

Commit 887c41c

Browse files
ekosmanfmassa
authored andcommitted
Expose num_workers in VideoClips (#1359)
1 parent 9ec6989 commit 887c41c

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

torchvision/datasets/video_utils.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,8 +49,9 @@ class VideoClips(object):
4949
on the resampled video
5050
"""
5151
def __init__(self, video_paths, clip_length_in_frames=16, frames_between_clips=1,
52-
frame_rate=None, _precomputed_metadata=None):
52+
frame_rate=None, _precomputed_metadata=None, num_workers=1):
5353
self.video_paths = video_paths
54+
self.num_workers = num_workers
5455
if _precomputed_metadata is None:
5556
self._compute_frame_pts()
5657
else:
@@ -77,7 +78,7 @@ def __getitem__(self, idx):
7778
dl = torch.utils.data.DataLoader(
7879
DS(self.video_paths),
7980
batch_size=16,
80-
num_workers=torch.get_num_threads(),
81+
num_workers=self.num_workers,
8182
collate_fn=lambda x: x)
8283

8384
with tqdm(total=len(dl)) as pbar:

0 commit comments

Comments
 (0)