forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
dataloader.pyi
39 lines (33 loc) · 1.72 KB
/
dataloader.pyi
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
from typing import Any, Callable, TypeVar, Generic, overload, Sequence, List
from . import Dataset, Sampler
T_co = TypeVar('T_co', covariant=True)
T = TypeVar('T')
_worker_init_fn_t = Callable[[int], None]
# Ideally we would parameterize `DataLoader` by the return type of `collate_fn`, but there is currently no way to have that
# type parameter set to a default value if the user doesn't pass in a custom 'collate_fn'.
# See https://github.com/python/mypy/issues/3737.
_collate_fn_t = Callable[[List[T]], Any]
class DataLoader(Generic[T_co]):
dataset: Dataset[T_co]
batch_size: int
num_workers: int
pin_memory: bool
drop_last: bool
timeout: float
@overload
def __init__(self, dataset: Dataset[T_co], batch_size: int=..., shuffle: bool=..., sampler: Sampler[int]=...,
num_workers: int=..., collate_fn: _collate_fn_t=..., pin_memory: bool=...,
drop_last: bool=..., timeout: float=..., worker_init_fn: _worker_init_fn_t=...) -> None: ...
@overload
def __init__(self, dataset: Dataset[T_co], batch_sampler: Sampler[Sequence[int]]=..., num_workers: int=...,
collate_fn: _collate_fn_t=..., pin_memory: bool=..., timeout: float=...,
worker_init_fn: _worker_init_fn_t=...) -> None: ...
def __len__(self) -> int: ...
# We quote '_DataLoaderIter' since it isn't defined yet and the definition can't be moved up since
# '_DataLoaderIter' references 'DataLoader'. Pending updates of PEP 484 will fix this.
def __iter__(self) -> '_DataLoaderIter':...
class _DataLoaderIter:
def __init__(self, loader: DataLoader) -> None:...
def __len__(self) -> int: ...
def __iter__(self) -> _DataLoaderIter: ...
def __next__(self) -> Any: ...