- Update BART settings

jalammar · Jan 9, 2022 · 58620b3 · 58620b3
1 parent 17bebb3
commit 58620b3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/ecco/lm.py b/src/ecco/lm.py
@@ -589,7 +589,7 @@ def _inhibit_neurons_hook(self, name: str, input_tensor):
         of the neurons indicated in self.neurons_to_inhibit
         """
 
-        layer_number = int(name.split('.')[2])
+        layer_number = re.search("(?<=\.)\d+(?=\.)", name).group(0)
         if layer_number in self.neurons_to_inhibit.keys():
             # print('layer_number', layer_number, input_tensor[0].shape)