fix an edge case bug of "get_layer"

ain-soph · Jul 13, 2022 · 356dc3d · 356dc3d
1 parent 67fbbcc
commit 356dc3d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/trojanzoo/utils/model.py b/trojanzoo/utils/model.py
@@ -356,7 +356,7 @@ def _get_layer(module: nn.Module, x: torch.Tensor,
                 layer_input = 'record'
             elif layer_input == full_name:
                 layer_input = 'record'
-            if layer_output.startswith(full_name):
+            if layer_output == full_name or layer_output.startswith(full_name + '.'):
                 return x
     else:
         x = module(x)