def save(self, path: str=None):
save_items = {"params": self.flow.params,
"state": self.flow.state,
"opt_state": self.trainer.opt_state,
"train_losses": self.train_losses,
"test_losses": self.test_losses}
util.save_pytree(save_items, path, overwrite=True)
def save(self, path: str=None):
save_items = {"params": self.params,
"state": self.state,
"test_eval_times": self.test_eval_times}
train_items = self.trainer.save_items()
test_items = self.tester.save_items()
save_items.update(train_items)
save_items.update(test_items)
util.save_pytree(save_items, path, overwrite=True)
def save(self, path: str=None):
save_items = {"params": self.params,
"state": self.state}
util.save_pytree(save_items, path, overwrite=True)