(self)
| 113 | return NetworkOutput(value, value_prefix, actor_logit, state, reward_hidden) |
| 114 | |
| 115 | def get_weights(self): |
| 116 | return {k: v.cpu() for k, v in self.state_dict().items()} |
| 117 | |
| 118 | def set_weights(self, weights): |
| 119 | self.load_state_dict(weights) |