Skip to content

Commit d886c89

Browse files
committed
Update deprecated arguments in torchvision
Signed-off-by: Beat Buesser <[email protected]>
1 parent a5f2759 commit d886c89

File tree

1 file changed

+6
-1
lines changed

1 file changed

+6
-1
lines changed

art/estimators/object_detection/pytorch_faster_rcnn.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -93,10 +93,15 @@ def __init__(
9393
if available otherwise run on CPU.
9494
"""
9595
import torchvision
96+
from torchvision.models import ResNet50_Weights
97+
from torchvision.models.detection import FasterRCNN_ResNet50_FPN_Weights
9698

9799
if model is None: # pragma: no cover
98100
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(
99-
pretrained=True, progress=True, num_classes=91, pretrained_backbone=True
101+
weights=FasterRCNN_ResNet50_FPN_Weights.COCO_V1,
102+
progress=True,
103+
num_classes=91,
104+
weights_backbone=ResNet50_Weights.IMAGENET1K_V1,
100105
)
101106

102107
super().__init__(

0 commit comments

Comments
 (0)