(self, hidden_states)
| 54 | self.flatten = nn.Flatten() |
| 55 | |
| 56 | def forward(self, hidden_states): |
| 57 | |
| 58 | if hidden_states.device != self.summary.weight.device: |
| 59 | hidden_states = hidden_states.to(self.summary.weight.device) |
| 60 | |
| 61 | output = self.dropout(hidden_states) |
| 62 | |
| 63 | # For now force upcast in fp32 if needed. Let's keep the |
| 64 | # output in fp32 for numerical stability. |
| 65 | if output.dtype != self.summary.weight.dtype: |
| 66 | output = output.to(self.summary.weight.dtype) |
| 67 | |
| 68 | output = self.summary(output) |
| 69 | values = torch.tanh(output).squeeze(-1) |
| 70 | return values |
| 71 | |
| 72 | |
| 73 | class AutoModelForCausalLMWithValueHead(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected