diff --git a/qa/python_models/torchvision/resnet50/model.py b/qa/python_models/torchvision/resnet50/model.py index 6a31a77420..46f83bcafd 100644 --- a/qa/python_models/torchvision/resnet50/model.py +++ b/qa/python_models/torchvision/resnet50/model.py @@ -1,4 +1,4 @@ -# Copyright 2023-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions @@ -25,6 +25,7 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. import torch +import torchvision import triton_python_backend_utils as pb_utils from torch.utils.dlpack import to_dlpack @@ -35,14 +36,9 @@ def initialize(self, args): This function initializes pre-trained ResNet50 model. """ self.device = "cuda" if args["model_instance_kind"] == "GPU" else "cpu" - # Avoid the "HTTP Error 403: rate limit exceeded" error - torch.hub._validate_not_a_forked_repo = lambda a, b, c: True - # Our tests currently depend on torchvision=0.14, - # to make sure `torch.hub` loads Resnet50 implementation - # compatible with torchvision=0.14, we need to provide tag self.model = ( - torch.hub.load( - "pytorch/vision:v0.14.1", "resnet50", weights="IMAGENET1K_V2" + torchvision.models.resnet50( + weights=torchvision.models.ResNet50_Weights.DEFAULT ) .to(self.device) .eval()