github
/
InternLM
mirror of https://github.com/InternLM/InternLM
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

fix pylint

pull/545/head
zigzagcai 2023-12-15 19:58:22 +08:00
parent 2dc8ddd582
commit 2afeebe5b0
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
internlm/train/training_internlm.py
View File

@ -263,8 +263,8 @@ def get_train_data_loader(num_worker: int = 0, dataset_generate_func: Optional[C
data_world_size=gpc.get_world_size(ParallelMode.DATA), data_world_size=gpc.get_world_size(ParallelMode.DATA),
) )
train_collate_fn = partial(packed_collate_fn, packed_length=data_cfg.packed_length) train_collate_fn = partial(packed_collate_fn, packed_length=data_cfg.packed_length)
def dl_worker_init(worker_id): def dl_worker_init(worker_id): # pylint: disable=unused-argument
gc.enable() gc.enable()
# Create the training data loader # Create the training data loader
@ -275,7 +275,7 @@ def get_train_data_loader(num_worker: int = 0, dataset_generate_func: Optional[C
pin_memory=True, pin_memory=True,
collate_fn=train_collate_fn, collate_fn=train_collate_fn,
persistent_workers=num_worker > 0, persistent_workers=num_worker > 0,
worker_init_fn = dl_worker_init, worker_init_fn=dl_worker_init,
) )
return train_dl, dataset_types return train_dl, dataset_types