Allow for attention weights to be extracted.
Browse filesThere is a small bug in indexing which didn't allow me to get the attentions out. I think this change fixes it.
- modeling_codesage.py +1 -1
modeling_codesage.py
CHANGED
|
@@ -149,7 +149,7 @@ class CodeSageBlock(nn.Module):
|
|
| 149 |
feed_forward_hidden_states = self.mlp(hidden_states)
|
| 150 |
hidden_states = residual + feed_forward_hidden_states
|
| 151 |
|
| 152 |
-
outputs = (hidden_states,) + outputs
|
| 153 |
return outputs # hidden_states, present, (attentions)
|
| 154 |
|
| 155 |
|
|
|
|
| 149 |
feed_forward_hidden_states = self.mlp(hidden_states)
|
| 150 |
hidden_states = residual + feed_forward_hidden_states
|
| 151 |
|
| 152 |
+
outputs = (hidden_states,) + outputs
|
| 153 |
return outputs # hidden_states, present, (attentions)
|
| 154 |
|
| 155 |
|