Spaces:
Running on Zero
Running on Zero
Vansh Chugh commited on
Commit ·
828d286
1
Parent(s): 6f9efe3
shorten comment
Browse files
src/audiocraft/modules/transformer.py
CHANGED
|
@@ -399,9 +399,8 @@ class StreamingMultiheadAttention(StreamingModule):
|
|
| 399 |
v = expand_repeated_kv(v, self.kv_repeat, self.memory_efficient)
|
| 400 |
if self.attention_as_float32:
|
| 401 |
q, k, v = [x.float() for x in [q, k, v]]
|
| 402 |
-
#
|
| 403 |
-
#
|
| 404 |
-
# align8-padding trim -- it needs the same seq_len match as the mask.
|
| 405 |
if custom_attn_mask:
|
| 406 |
# When using a custom attn mask:
|
| 407 |
# Move to query's device, repeat for each sample, remove align8 padding
|
|
|
|
| 399 |
v = expand_repeated_kv(v, self.kv_repeat, self.memory_efficient)
|
| 400 |
if self.attention_as_float32:
|
| 401 |
q, k, v = [x.float() for x in [q, k, v]]
|
| 402 |
+
# block was initially reachable inside 'if self.memory_efficient'
|
| 403 |
+
# moved out so plain attn also gets trimmed to match seq_len's mask.
|
|
|
|
| 404 |
if custom_attn_mask:
|
| 405 |
# When using a custom attn mask:
|
| 406 |
# Move to query's device, repeat for each sample, remove align8 padding
|