def METHOD_NAME(self):
    model = ModelForTest()
    with TemporaryDirectory() as tmp_dir:
        offload_state_dict(tmp_dir, model.state_dict())
        index_file = os.path.join(tmp_dir, "index.json")
        self.assertTrue(os.path.isfile(index_file))
        # TODO: add tests on what is inside the index
        for key in ["linear1.weight", "linear1.bias", "linear2.weight", "linear2.bias"]:
            weight_file = os.path.join(tmp_dir, f"{key}.dat")
            self.assertTrue(os.path.isfile(weight_file))
            # TODO: add tests on the fact weights are properly loaded