From ac1450f4fcd5d7f6c36d721ef62f164ab1cd1ae2 Mon Sep 17 00:00:00 2001
From: Hao Zhang <hzhangxyz@outlook.com>
Date: Fri, 27 Mar 2026 10:01:45 +0800
Subject: [PATCH] fix: detach k, v in SelfAttention to prevent gradient
 accumulation

Return detached key and value tensors to avoid unintended gradient
tracking during backpropagation.
---
 qmp/networks/transformers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/qmp/networks/transformers.py b/qmp/networks/transformers.py
index 3ddba315..47117168 100644
--- a/qmp/networks/transformers.py
+++ b/qmp/networks/transformers.py
@@ -81,7 +81,7 @@ def forward(
         # attn: batch, heads_num, site, heads_dim
         out = attn.transpose(1, 2).reshape([batch_size, sites, embedding_dim])
         # out: batch, site, embedding_dim
-        return self.out(out), (k, v)
+        return self.out(out), (k.detach(), v.detach())
 
 
 class DecoderUnit(torch.nn.Module):