yolov11_heatmap.py not working for classification

[Mordokkai]

Dear @z1069614715 ,

I'm trying to use your script yolov11_heatmap.py on yolo11-cls.pt for classification task. However I'm getting the following error:
Warning... self.method(tensor, [self.target]) failure.
'tuple' object has no attribute 'cpu'

My params are:

def get_params():
    params = {
        'device': 'cuda:0',
        'method': 'GradCAM',
        'layer': [8, 10], #, 12, 14, 16, 18],
        'backward_type': "classify", 
        'conf_threshold': 0.2, 
        'ratio': 0.02, 
        'show_result': True, 
        'renormalize': False, 
        'task':'classify', 
        'img_size':640, 
    }

By the way, for detection task it's working.

Could you help me on this issue ?
Bests,