def test__training_step__log_max_reduce_fx(tmp_path, batches, fx, result):
    """Tests that log works correctly with different tensor types."""

    class TestModel(BoringModel):
        def training_step(self, batch, batch_idx):
            acc = self.step(batch[0])
            self.log("foo", torch.tensor(batch_idx, dtype=torch.long), on_step=False, on_epoch=True, reduce_fx=fx)
            return acc

        def validation_step(self, batch, batch_idx):
            loss = self.step(batch)
            self.log("bar", torch.tensor(batch_idx).float(), on_step=False, on_epoch=True, reduce_fx=fx)
            return {"x": loss}

    model = TestModel()
    trainer = Trainer(
        default_root_dir=tmp_path,
        limit_train_batches=batches,
        limit_val_batches=batches,
        max_epochs=2,
        enable_model_summary=False,
    )
    trainer.fit(model)

    # make sure types are correct
    assert trainer.logged_metrics["foo"] == result
    assert trainer.logged_metrics["bar"] == result
Tests that log works correctly with different tensor types.
tests/tests_pytorch/trainer/logging_/test_train_loop_logging.py