add reset

jacobgil · Dec 31, 2020 · 1d68f22 · 1d68f22
1 parent 4f75bc6
commit 1d68f22
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/vit_rollout.py b/vit_rollout.py
@@ -57,6 +57,7 @@ def get_attention(self, module, input, output):
         self.attentions.append(output.cpu())
 
     def __call__(self, input_tensor):
+        self.attentions = []
         with torch.no_grad():
             output = self.model(input_tensor)