# Copyright (c) OpenMMLab. All rights reserved.
import torch
from mmcls.models import ImageClassifier

from mmrazor.registry import TASK_UTILS


@TASK_UTILS.register_module()
class ImageClassifierPseudoLoss:
    """Calculate the pseudo loss to trace the topology of a `ImageClassifier`
    in MMClassification with `BackwardTracer`."""

    def __call__(self, model: ImageClassifier) -> torch.Tensor:
        pseudo_img = torch.rand(1, 3, 224, 224)
        pseudo_output = model(pseudo_img)
        return sum(pseudo_output)
