fix: no pinning memory on CPU (#4874)

caic99 · pre-commit-ci[bot] · web-flow · commit 3be075569ccb · 2025-08-07T20:44:23.000Z
This PR fixes the warning message when running on CPU:
```
torch/utils/data/dataloader.py:665: UserWarning: 'pin_memory' argument is set as true but no accelerator is found, then device pinned memory won't be used.
  warnings.warn(warn_msg)
  ```

&lt;!-- This is an auto-generated comment: release notes by coderabbit.ai --&gt;

## Summary by CodeRabbit

* **Bug Fixes**
  * Improved memory handling by enabling memory pinning only when using non-CPU devices, which may enhance performance and stability when training on GPUs.

&lt;!-- end of auto-generated comment: release notes by coderabbit.ai --&gt;

---------

Co-authored-by: pre-commit-ci[bot] &lt;66853113+pre-commit-ci[bot]@users.noreply.github.com&gt;
diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
@@ -195,7 +195,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
                 )
                 _data_iter = cycle_iterator(_dataloader)
                 return _dataloader, _data_iter

Original file line number	Diff line number	Diff line change
`@@ -195,7 +195,7 @@ def get_dataloader_and_iter(_data, _params):`
`195`	`195`	`else 0, # setting to 0 diverges the behavior of its iterator; should be >=1`
`196`	`196`	`drop_last=False,`
`197`	`197`	`collate_fn=lambda batch: batch, # prevent extra conversion`
`198`		`- pin_memory=True,`
	`198`	`+ pin_memory=(DEVICE != "cpu"), # pin memory only if not on CPU`
`199`	`199`	`)`
`200`	`200`	`_data_iter = cycle_iterator(_dataloader)`
`201`	`201`	`return _dataloader, _data_iter`