sliding_window_inference not working with torch create_feature_extractor for attention layers

[CataRay]

I am trying to get the attention layers out of a model I have trained.
I can evaluate correctly the model by using sliding_window_inference on the entire 3D volume. But when I try to use the same function to get the output of those layers I get the error:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Input In [18], in <cell line: 18>()
     21 sw_batch_size = 4
     22 test_data["pred"] = sliding_window_inference(test_inputs, roi_size, sw_batch_size, trained_model)
---> 23 test_data["att2"] = sliding_window_inference(test_inputs, roi_size, sw_batch_size, trained_model_ext2)
     24 test_data["att3"] = sliding_window_inference(test_inputs, roi_size, sw_batch_size, trained_model_ext3)
     25 test_data["att4"] = sliding_window_inference(test_inputs, roi_size, sw_batch_size, trained_model_ext4)

File /trials/Catalina/02_Projects/01_DeepLearning/_virtual_envs_/venv_Monai/lib/python3.8/site-packages/monai/inferers/utils.py:130, in sliding_window_inference(inputs, roi_size, sw_batch_size, predictor, overlap, mode, sigma_scale, padding_mode, cval, sw_device, device, *args, **kwargs)
    125 unravel_slice = [
    126     [slice(int(idx / num_win), int(idx / num_win) + 1), slice(None)] + list(slices[idx % num_win])
    127     for idx in slice_range
    128 ]
    129 window_data = torch.cat([inputs[win_slice] for win_slice in unravel_slice]).to(sw_device)
--> 130 seg_prob = predictor(window_data, *args, **kwargs).to(device)  # batched patch segmentation
    132 if not _initialized:  # init. buffer at the first iteration
    133     output_classes = seg_prob.shape[1]

AttributeError: 'dict' object has no attribute 'to'

I am not sure if I am missing something here but I can't find much information on this topic so far. Here is the eval code:

from tqdm.auto import tqdm
from monai.inferers import sliding_window_inference
from torchvision.models.feature_extraction import create_feature_extractor

#extract last attention layer
if att:
    trained_model_ext2 = create_feature_extractor(trained_model, {'Att2.psi': 'att2'})
    trained_model_ext3 = create_feature_extractor(trained_model, {'Att3.psi': 'att3'})
    trained_model_ext4 = create_feature_extractor(trained_model, {'Att4.psi': 'att4'})
    trained_model_ext5 = create_feature_extractor(trained_model, {'Att5.psi': 'att5'})
    trained_model_ext2.eval()
    trained_model_ext3.eval()
    trained_model_ext4.eval()
    trained_model_ext5.eval()

with torch.no_grad():
    for i, test_data in enumerate(test_dataloader):
        test_inputs = test_data["input_images"].to(device)
        sw_batch_size = 4
        test_data["pred"] = sliding_window_inference(test_inputs, roi_size, sw_batch_size, trained_model)
        test_data["att2"] = sliding_window_inference(test_inputs, roi_size, sw_batch_size, trained_model_ext2)
        test_data["att3"] = sliding_window_inference(test_inputs, roi_size, sw_batch_size, trained_model_ext3)
        test_data["att4"] = sliding_window_inference(test_inputs, roi_size, sw_batch_size, trained_model_ext4)
        test_data["att5"] = sliding_window_inference(test_inputs, roi_size, sw_batch_size, trained_model_ext5)

[wyli]

I'm not sure what kind of trained_model you have, but monai v0.9 released a few days ago now supports predictors with dictionary outputs, perhaps worth trying...