Используя 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 использовать только выходные данные сегментации из модели с несколькими выходами?
Обновление:
Могу ли я определить функцию-член в классе модели:

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

def predict_seg(self, x: torch.Tensor) -> torch.Tensor:
out = self.forward(x)
return out["seg_logits"]

Затем класс-оболочка
class SegmentationWrapper(nn.Module):
def init(self, model):
super().init()
self.model = model

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

def forward(self, x):
return self.model.predict_seg(x)

тогда:

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

model = SegmentationWrapper(model)

Затем вызываем MONAI:

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

predictor=model

Это правильное решение? Любая помощь приветствуется. Заранее спасибо!