PyTorch torchtune

# レシピファイルの `train.py` 関数内
self._metric_logger.log_dict(
    {"current_epoch": self.epochs * self.global_step / self._steps_per_epoch},
    step=self.global_step,
)

def save_checkpoint(self, epoch: int) -> None:
    ...
    ## W&B にチェックポイントを保存します。
    ## ファイル名は Checkpointer クラスによって異なります。
    ## 以下は full_finetune の場合の例です。
    checkpoint_file = Path.joinpath(
        self._checkpointer._output_dir, f"torchtune_model_{epoch}"
    ).with_suffix(".pt")
    wandb_artifact = wandb.Artifact(
        name=f"torchtune_model_{epoch}",
        type="model",
        # モデル チェックポイントの説明
        description="Model checkpoint",
        # 任意のメタデータを dict として追加できます
        metadata={
            utils.SEED_KEY: self.seed,
            utils.EPOCHS_KEY: self.epochs_run,
            utils.TOTAL_EPOCHS_KEY: self.total_epochs,
            utils.MAX_STEPS_KEY: self.max_steps_per_epoch,
        },
    )
    wandb_artifact.add_file(checkpoint_file)
    wandb.log_artifact(wandb_artifact)