Transformers
PyTorch
code
custom_code
Inference Endpoints

Allow for attention weights to be extracted.

#2
by FJFehr - opened
Files changed (1) hide show
  1. 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[1:]
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