Fix requires_grad for identity in TNet

adosar · Dec 2, 2024 · 5e8bda5 · 5e8bda5
1 parent fed7f75
commit 5e8bda5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/aidsorb/modules.py b/src/aidsorb/modules.py
@@ -207,7 +207,7 @@ def forward(self, x):
         x = self.dense_blocks(x)
 
         # Initialize the identity matrix.
-        identity = torch.eye(self.embed_dim, device=x.device, requires_grad=True).repeat(bs, 1, 1)
+        identity = torch.eye(self.embed_dim, device=x.device, requires_grad=x.requires_grad).repeat(bs, 1, 1)
 
         # Output has shape (B, self.embed_dim, self.embed_dim).
         x = x.view(-1, self.embed_dim, self.embed_dim) + identity