(self, hidden_state, with_grad=True)
| 580 | return reward_w_dist, representation_mean, dynamic_mean, reward_mean |
| 581 | |
| 582 | def project(self, hidden_state, with_grad=True): |
| 583 | # only the branch of proj + pred can share the gradients |
| 584 | hidden_state = hidden_state.view(-1, self.porjection_in_dim) |
| 585 | proj = self.projection(hidden_state) |
| 586 | |
| 587 | # with grad, use proj_head |
| 588 | if with_grad: |
| 589 | proj = self.projection_head(proj) |
| 590 | return proj |
| 591 | else: |
| 592 | return proj.detach() |
| 593 |