копировать сегмент из тензора изображения

0

У меня есть три тензора: A - (1, 3, 256, 256) B - (1, 3, 256, 256) - это тензор белого изображения C - (256, 256) - это тензор сегментов

Например, C будет выглядеть так:

tensor([[ 337,  337,  337,  ...,  340,  340,  340],
        [ 337,  337,  337,  ...,  340,  340,  340],
        [ 337,  337,  337,  ...,  340,  340,  340],
        ...,
        [1022, 1022, 1022,  ..., 1010, 1010, 1010],
        [1022, 1022, 1022,  ..., 1010, 1010, 1010],
        [1022, 1022, 1022,  ..., 1010, 1010, 1010]], device='cuda:0')

где 37 может указывать на здание и т. д.

Тензор C определяет расположение формы сегмента. Я хочу скопировать один и тот же сегмент на основе местоположения из тензора A на тензор B. Это будет фотошоп сегмента на тензор белого изображения.

Это похоже на маскирование, и я просмотрел mask_select ( https://pytorch.org/docs/stable/generated/torch.masked_select.html ), но он возвращает только 1D тензор.