solve issue with messages (#35)

facebookresearch · Jun 17, 2024 · 391ae43 · 391ae43
2 parents 673c000 + da8a29c
commit 391ae43
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/src/audioseal/models.py b/src/audioseal/models.py
@@ -109,13 +109,13 @@ def get_watermark(
         hidden = self.encoder(x)
 
         if self.msg_processor is not None:
-            if self.message is None:
-                self.message = torch.randint(0, 2, (x.shape[0], self.msg_processor.nbits), device=x.device)
+            if message is None:
+                if self.message is None:
+                    message = torch.randint(0, 2, (x.shape[0], self.msg_processor.nbits), device=x.device)
+                else:
+                    message = self.message.to(device=x.device)
             else:
-                self.message = self.message.to(device=x.device)
-
-
-            message = self.message
+                message = message.to(device=x.device)
 
             hidden = self.msg_processor(hidden, message)