Merge pull request #142 from rahulxie/main

fix a bug in deep_inspect
ain-soph · May 26, 2022 · cc4551f · cc4551f
2 parents fb68f5f + 22059ba
commit cc4551f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/trojanvision/defenses/backdoor/model_inspection/deep_inspect.py b/trojanvision/defenses/backdoor/model_inspection/deep_inspect.py
@@ -45,10 +45,11 @@ def __init__(self, defense_remask_epoch: int = 20, defense_remask_lr=0.01,
         subset, _ = self.dataset.split_dataset(dataset, percent=sample_ratio)
         self.loader = self.dataset.get_dataloader(mode='train', dataset=subset)
 
-    def optimize_mark(self, label: int) -> tuple[torch.Tensor, float]:
+    def optimize_mark(self, label: int, **kwargs) -> tuple[torch.Tensor, float]:
         r"""
         Args:
             label (int): The class label to optimize.
+            **kwargs: Any keyword argument (unused).
 
         Returns:
             (torch.Tensor, torch.Tensor):