fix: no pinning memory on CPU

caic99 · caic99 · commit c63bfe94d68f · 2025-08-07T19:21:09.000+08:00
diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
@@ -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
                 )
                 _data_iter = cycle_iterator(_dataloader)
                 return _dataloader, _data_iter

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