omer11a commited on
Commit
1127ecd
1 Parent(s): 3342cb9

Improved memory efficiency

Browse files
Files changed (1) hide show
  1. bounded_attention.py +1 -1
bounded_attention.py CHANGED
@@ -301,7 +301,7 @@ class BoundedAttention(injection_utils.AttentionBase):
301
 
302
  mask_maps.append(attn.mean(dim=1)) # mean over heads
303
  if self.cur_step > 0:
304
- mask_maps = mask_maps[1:] # throw away old maps
305
 
306
  def _save_loss_values(self, attn, is_cross):
307
  if (
 
301
 
302
  mask_maps.append(attn.mean(dim=1)) # mean over heads
303
  if self.cur_step > 0:
304
+ mask_maps.pop(0) # throw away old maps
305
 
306
  def _save_loss_values(self, attn, is_cross):
307
  if (