fix

ain-soph · Dec 25, 2023 · 3480438 · 3480438
1 parent 1e11584
commit 3480438
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/trojanvision/attacks/backdoor/normal/imc.py b/trojanvision/attacks/backdoor/normal/imc.py
@@ -66,7 +66,7 @@ def optimize_mark(self, loss_fn: Callable[..., torch.Tensor] = None, **kwargs):
         r"""Optimize watermark at the beginning of each training epoch."""
         loss_fn = loss_fn or self.model.loss
 
-        atanh_mark = torch.randn_like(self.mark.mark[:-1], requires_grad=True)
+        atanh_mark = self.mark.mark[:-1].mul(2).sub(1).mul(0.999).atanh().detach().requires_grad_()
         optimizer = optim.Adam([atanh_mark], lr=self.attack_remask_lr)
         optimizer.zero_grad()