Repository: Lightning-AI/pytorch-lightning
File: tests/metrics/classification/test_f_beta.py (Python)

Code before:
            _multidim_multiclass_prob_inputs.target,
            _sk_fbeta_multidim_multiclass_prob,
            NUM_CLASSES,
            False,
        ),
        (
            _multidim_multiclass_inputs.preds,
            _multidim_multiclass_inputs.target,
            _sk_fbeta_multidim_multiclass,
            NUM_CLASSES,
            False,
        ),
    ],
)
@pytest.mark.parametrize("average", ['micro', 'macro', 'weighted', None])
@pytest.mark.parametrize("beta", [0.5, 1.0, 2.0])
class TestFBeta(MetricTester):
    @pytest.mark.parametrize("ddp", [True, False])
    @pytest.mark.parametrize("dist_sync_on_step", [True, False])
    def test_fbeta(
        self, preds, target, sk_metric, num_classes, multilabel, average, beta, ddp, dist_sync_on_step
    ):
        if beta == 1.0:
            metric_class = F1
        else:
            metric_class = partial(FBeta, beta=beta)

        self.run_class_metric_test(
            ddp=ddp,
            preds=preds,
            target=target,
            metric_class=metric_class,
            sk_metric=partial(sk_metric, average=average, beta=beta),
            dist_sync_on_step=dist_sync_on_step,
            metric_args={
                "num_classes": num_classes,
                "average": average,
                "multilabel": multilabel,
                "threshold": THRESHOLD,
            },
            check_dist_sync_on_step=False,
            check_batch=False,
        )

    def test_fbeta_functional(
        self, preds, target, sk_metric, num_classes, multilabel, average, beta
    ):
        if beta == 1.0:
            metric_functional = f1
        else:
            metric_functional = partial(fbeta, beta=beta)

Review comment (suggestion):
```suggestion
        metric_class = F1 if beta == 1.0 else partial(FBeta, beta=beta)
```

Code after:
            _multidim_multiclass_prob_inputs.target,
            _sk_fbeta_multidim_multiclass_prob,
            NUM_CLASSES,
            False,
        ),
        (
            _multidim_multiclass_inputs.preds,
            _multidim_multiclass_inputs.target,
            _sk_fbeta_multidim_multiclass,
            NUM_CLASSES,
            False,
        ),
    ],
)
@pytest.mark.parametrize("average", ['micro', 'macro', 'weighted', None])
@pytest.mark.parametrize("beta", [0.5, 1.0, 2.0])
class TestFBeta(MetricTester):
    @pytest.mark.parametrize("ddp", [True, False])
    @pytest.mark.parametrize("dist_sync_on_step", [True, False])
    def test_fbeta(
        self, preds, target, sk_metric, num_classes, multilabel, average, beta, ddp, dist_sync_on_step
    ):
        metric_class = F1 if beta == 1.0 else partial(FBeta, beta=beta)

        self.run_class_metric_test(
            ddp=ddp,
            preds=preds,
            target=target,
            metric_class=metric_class,
            sk_metric=partial(sk_metric, average=average, beta=beta),
            dist_sync_on_step=dist_sync_on_step,
            metric_args={
                "num_classes": num_classes,
                "average": average,
                "multilabel": multilabel,
                "threshold": THRESHOLD,
            },
            check_dist_sync_on_step=False,
            check_batch=False,
        )

    def test_fbeta_functional(
        self, preds, target, sk_metric, num_classes, multilabel, average, beta
    ):
        metric_functional = f1 if beta == 1.0 else partial(fbeta, beta=beta)

        self.run_functional_metric_test(preds=preds,
                                        target=target,
                                        metric_functional=metric_functional,
                                        sk_metric=partial(sk_metric, average=average, beta=beta),
                                        metric_args={