Skip to content

Commit d040d08

Browse files
wanghan-iapcmHan Wang
andauthored
fix: data type of nloc, nall-nloc in the input of border_op (#4653)
<!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **Refactor** - Updated internal processing to explicitly assign certain computational operations to run on the CPU. This change clarifies execution behavior and enhances consistency in data handling while preserving all existing functionalities for end-users. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>
1 parent 5831505 commit d040d08

File tree

2 files changed

+8
-8
lines changed

2 files changed

+8
-8
lines changed

deepmd/pt/model/descriptor/repflows.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -505,13 +505,13 @@ def forward(
505505
torch.tensor(
506506
real_nloc,
507507
dtype=torch.int32,
508-
device=env.DEVICE,
509-
), # should be int of c++
508+
device=torch.device("cpu"),
509+
), # should be int of c++, placed on cpu
510510
torch.tensor(
511511
real_nall - real_nloc,
512512
dtype=torch.int32,
513-
device=env.DEVICE,
514-
), # should be int of c++
513+
device=torch.device("cpu"),
514+
), # should be int of c++, placed on cpu
515515
)
516516
node_ebd_ext = ret[0].unsqueeze(0)
517517
if has_spin:

deepmd/pt/model/descriptor/repformers.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -491,13 +491,13 @@ def forward(
491491
torch.tensor(
492492
real_nloc,
493493
dtype=torch.int32,
494-
device=env.DEVICE,
495-
), # should be int of c++
494+
device=torch.device("cpu"),
495+
), # should be int of c++, placed on cpu
496496
torch.tensor(
497497
real_nall - real_nloc,
498498
dtype=torch.int32,
499-
device=env.DEVICE,
500-
), # should be int of c++
499+
device=torch.device("cpu"),
500+
), # should be int of c++, placed on cpu
501501
)
502502
g1_ext = ret[0].unsqueeze(0)
503503
if has_spin:

0 commit comments

Comments
 (0)