[Models] Improve iteration over layers (#19497)
Signed-off-by: Lukas Geiger <lukas.geiger94@gmail.com>
This commit is contained in:
@@ -345,8 +345,7 @@ class BambaModel(nn.Module):
|
||||
|
||||
residual = None
|
||||
num_attn = 0
|
||||
for i in range(len(self.layers)):
|
||||
layer = self.layers[i]
|
||||
for i, layer in enumerate(self.layers):
|
||||
if isinstance(layer, BambaAttentionDecoderLayer):
|
||||
num_attn += 1
|
||||
|
||||
|
||||
Reference in New Issue
Block a user