"src/vscode:/vscode.git/clone" did not exist on "5e8e6cb44f78fc9235455d08a79010425c9e5a24"
Commit cd9c320d authored by Naveen Suda's avatar Naveen Suda Committed by Facebook GitHub Bot
Browse files

add collate_fn to dataloader

Summary: Pull Request resolved: https://github.com/facebookresearch/d2go/pull/483

Reviewed By: YXIE14

Differential Revision: D42733542

fbshipit-source-id: 0dc936c536554b5beead462eaf74bc007758c12e
parent 3111ae59
......@@ -208,7 +208,7 @@ class D2GoDataAPIMixIn:
@classmethod
def build_detection_test_loader(
cls, cfg, dataset_name: Union[str, List[str]], mapper=None
cls, cfg, dataset_name: Union[str, List[str]], mapper=None, collate_fn=None
):
logger.info(
"Building detection test loader for dataset: {} ...".format(dataset_name)
......@@ -216,7 +216,9 @@ class D2GoDataAPIMixIn:
with configure_dataset_creation(cfg):
mapper = mapper or cls.get_mapper(cfg, is_train=False)
logger.info("Using dataset mapper:\n{}".format(mapper))
return d2_build_detection_test_loader(cfg, dataset_name, mapper=mapper)
return d2_build_detection_test_loader(
cfg, dataset_name, mapper=mapper, collate_fn=collate_fn
)
@classmethod
def build_detection_train_loader(cls, cfg, *args, mapper=None, **kwargs):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment