Class CustomModuleForSaveLoad

python/oneflow/test/modules/test_save_load.py:32–38

class CustomModuleForSaveLoad(flow.nn.Module):
    def __init__(self):
        super().__init__()
        self.param = flow.nn.Parameter(flow.randn(1, 3, 3, 3))

    def forward(self, x):
        return self.param + x


@unittest.skipIf(os.getenv("ONEFLOW_TEST_CPU_ONLY"), "only test cpu cases")