fix:device

anyangml · anyangml · commit 250168bbdc3b · 2026-03-25T14:42:37.000+08:00
diff --git a/deepmd/pt/loss/xas.py b/deepmd/pt/loss/xas.py
@@ -253,7 +253,9 @@ def forward(
             edge_idx = torch.zeros(nf, dtype=torch.long, device=pred.device)
 
         # e_ref_frame: [nf, 2]  (E_min_ref, E_max_ref for each frame)
-        e_ref_frame = self.e_ref[sel_type, edge_idx]  # [nf, 2]
+        # Indices must be on the same device as the buffer (handles CPU/GPU mismatch)
+        _dev = self.e_ref.device
+        e_ref_frame = self.e_ref[sel_type.to(_dev), edge_idx.to(_dev)].to(pred.device)
 
         # Shift the energy-dim TARGETS only.
         #

Original file line number	Diff line number	Diff line change
`@@ -253,7 +253,9 @@ def forward(`
`253`	`253`	`edge_idx = torch.zeros(nf, dtype=torch.long, device=pred.device)`
`254`	`254`
`255`	`255`	`# e_ref_frame: [nf, 2] (E_min_ref, E_max_ref for each frame)`
`256`		`- e_ref_frame = self.e_ref[sel_type, edge_idx] # [nf, 2]`
	`256`	`+ # Indices must be on the same device as the buffer (handles CPU/GPU mismatch)`
	`257`	`+ _dev = self.e_ref.device`
	`258`	`+ e_ref_frame = self.e_ref[sel_type.to(_dev), edge_idx.to(_dev)].to(pred.device)`
`257`	`259`
`258`	`260`	`# Shift the energy-dim TARGETS only.`
`259`	`261`	`#`