Error NonZero(/rpn/NonZero) when importing faster rcnn from torchvision.models

[JohnConnor123]

Hello. I tried to port the trained Faster RCNN from torchvision.models to RK3588 (orange pi 5), but I get an error in the photo. When trying to import a model written from scratch (article: https://towardsdatascience.com/understanding-and-implementing-faster-r-cnn-a-step-by-step-guide-11acfff216b0, github with code: https:/ /github.com/wingedrasengan927/pytorch-tutorials/tree/master/Object%20Detection) I get about the same error with the rpn layer. Importing libtorch didn't help. How can I import at least the standard, not fully trained FasterRCNN from torchvision.models? I've been struggling with this problem for over a week now.

My code:

model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 2)
big_model.load_state_dict(torch.load(big_model_path))
big_model.eval()