def test_comet_epoch_logging(comet_mock, tmp_path, monkeypatch):
    """Test that CometLogger removes the epoch key from the metrics dict and passes it as argument."""
    _patch_comet_atexit(monkeypatch)
    logger = CometLogger(project_name="test", save_dir=str(tmp_path))
    logger.log_metrics({"test": 1, "epoch": 1}, step=123)
    logger.experiment.__internal_api__log_metrics__.assert_called_once_with(
        {"test": 1},
        epoch=1,
        step=123,
        prefix=logger._prefix,
        framework="pytorch-lightning",
    )
Test that CometLogger removes the epoch key from the metrics dict and passes it as argument.
tests/tests_pytorch/loggers/test_comet.py