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
- # Original: this block was only reachable inside `if self.memory_efficient:`.
403
- # Moved out so the plain (non-memory-efficient) path below also gets the
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