Pytorch выделяет пиксели в ограничивающих прямоугольниках

Question

Код ниже получает маску для всех ограничивающих рамок, которые будут отображаться в opencv. Есть ли более быстрый способ сделать это только с pytorch или torchvision? Такое ощущение, что для этого должен быть один вкладыш.

import torch

boxes = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11,12]])
height = 720
width = 1280
noOfBoxes = 3
masklayersY = torch.arange(height).unsqueeze(1).unsqueeze(1).expand(-1, width, noOfBoxes)  # 720, 1280, 3
masklayersX = torch.arange(width).unsqueeze(1).unsqueeze(0).expand(height, -1, noOfBoxes)  # 720, 1280, 3

mask = (masklayersX >= boxes[:, 0]) & \
       (masklayersX <= boxes[:, 2]) & \
       (masklayersY >= boxes[:, 1]) & \
       (masklayersY <= boxes[:, 3])

maskT, _ = mask.max(dim=2)
maskT = (maskT * 255).to(torch.uint8)

Это можно сделать с помощью for для l oop над ячейками, но я думаю, это должно быть возможно векторизованным способом непосредственно в pytorch.