def test_module_current_fx_attributes_reset(tmp_path):
    """Ensure that lightning module's attributes related to current fx are reset at the end of execution."""
    model = BoringModel()
    trainer = Trainer(default_root_dir=tmp_path, fast_dev_run=1, enable_checkpointing=False, logger=False)

    trainer.fit(model)
    assert model._current_fx_name is None

    trainer.test(model)
    assert model._current_fx_name is None
Ensure that lightning module's attributes related to current fx are reset at the end of execution.
tests/tests_pytorch/trainer/test_trainer.py