"configs/convnext_v2/convnext-v2-pico_32xb32_in1k.py" did not exist on "59b09903af84245714720bfd329ab9f5f005ce90"
savable.py 614 Bytes
Newer Older
jerrrrry's avatar
jerrrrry committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
# Copyright (c) 2025, NVIDIA CORPORATION.
# SPDX-License-Identifier: BSD-3-Clause

from abc import ABC, abstractmethod

from megatron.energon.state import FlexState


class Savable(ABC):
    @abstractmethod
    def save_state(self) -> FlexState:
        """Dumps the state of the object into some serializable format.
        Should only dump the part that is persistable for energon usage.
        """
        ...

    @abstractmethod
    def restore_state(self, state: FlexState):
        """Restores the state of the object.
        If state is None, the object is restored to its default state."""
        ...