Skip to content
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion deepmd/pt/train/training.py
Original file line number Diff line number Diff line change
Expand Up @@ -194,7 +194,7 @@ def get_dataloader_and_iter(_data, _params):
else 0, # setting to 0 diverges the behavior of its iterator; should be >=1
drop_last=False,
collate_fn=lambda batch: batch, # prevent extra conversion
pin_memory=True,
pin_memory=(DEVICE != "cpu"), # pin memory only if not on CPU
Comment thread
njzjz marked this conversation as resolved.
Outdated
)
_data_iter = cycle_iterator(_dataloader)
return _dataloader, _data_iter
Expand Down
Loading