r"""Returns a dictionary containing whole states of the module.
(self, rst=None, prefix="", keep_var=False)
| 467 | return rst |
| 468 | |
| 469 | def _state_dict(self, rst=None, prefix="", keep_var=False): |
| 470 | r"""Returns a dictionary containing whole states of the module.""" |
| 471 | |
| 472 | def is_state(obj): |
| 473 | return _is_parameter(obj) or _is_buffer(obj) |
| 474 | |
| 475 | module_type = self.__class__ |
| 476 | if rst is None: |
| 477 | rst = OrderedDict() |
| 478 | |
| 479 | for k, v in self._flatten(recursive=False, with_key=True, predicate=is_state): |
| 480 | assert prefix + k not in rst, "duplicated state: {}".format(k) |
| 481 | if keep_var: |
| 482 | rst[(module_type, prefix + k)] = v |
| 483 | else: |
| 484 | rst[(module_type, prefix + k)] = v.numpy() |
| 485 | |
| 486 | for k, submodule in self._flatten( |
| 487 | recursive=False, |
| 488 | with_key=True, |
| 489 | predicate=lambda obj: isinstance(obj, Module), |
| 490 | ): |
| 491 | submodule.state_dict(rst, prefix + k + ".", keep_var) |
| 492 | |
| 493 | return rst |
| 494 | |
| 495 | def load_state_dict( |
| 496 | self, |
no test coverage detected