def save_attn_gradients(self, attn_gradients):
    self.attn_gradients = attn_gradients
    
def get_attn_gradients(self):
    return self.attn_gradients

def save_attn_map(self, attention_map):
    self.attention_map = attention_map
    
def get_attn_map(self):
    return self.attention_map