Используя MONAI moving_window_inference с моделью с несколькими выходами, как использовать предиктор только для одного в ⇐ Python

[Anonymous]

Я использую MONAI moving_window_inference для вывода сегментации, но моя модель возвращает несколько выходных данных вместо одного тензора.
Моя модель вперед возвращает что-то вроде:

Код: Выделить всё

return {
"latent_a": latent_a,
"segmentation": segmentation_logits,
"classification": classification_logits,
}

Для вывода скользящего окна я хочу использовать только вывод «сегментации».
Документы MONAI говорят, что аргумент предиктора может быть вызываемым, но я не уверен, какой самый чистый или рекомендуемый способ использовать его в моем случае.
В настоящее время я это делаю:

Код: Выделить всё

outputs = sliding_window_inference(
inputs=images,
roi_size=(256, 256),
sw_batch_size=1,
predictor=model,
)

Как правильно MONAI заставитьскользящий_window_inference использовать только выходные данные сегментации из модели с несколькими выходами?