Spaces:
Running
on
Zero
Running
on
Zero
Improved memory efficiency
Browse files- 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
|
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 (
|