fix(pt): detach computed descriptor tensor to prevent OOM

Fix #4544. Signed-off-by: Jinzhe Zeng <[email protected]>
deepmodeling · Jan 11, 2025 · 9f37532 · 9f37532
1 parent 65ca05a
commit 9f37532
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py
@@ -244,7 +244,7 @@ def forward_atomic(
         )
         assert descriptor is not None
         if self.enable_eval_descriptor_hook:
-            self.eval_descriptor_list.append(descriptor)
+            self.eval_descriptor_list.append(descriptor.detach())
         # energy, force
         fit_ret = self.fitting_net(
             descriptor,