Skip to content

Commit a161098

Browse files
authored
Update ShuffleNetV2 annotations for x1_5 and x2_0 variants (#6022)
1 parent 01664a8 commit a161098

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

torchvision/models/shufflenetv2.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -316,7 +316,7 @@ def shufflenet_v2_x1_0(
316316
return _shufflenetv2(weights, progress, [4, 8, 4], [24, 116, 232, 464, 1024], **kwargs)
317317

318318

319-
@handle_legacy_interface(weights=("pretrained", None))
319+
@handle_legacy_interface(weights=("pretrained", ShuffleNet_V2_X1_5_Weights.IMAGENET1K_V1))
320320
def shufflenet_v2_x1_5(
321321
*, weights: Optional[ShuffleNet_V2_X1_5_Weights] = None, progress: bool = True, **kwargs: Any
322322
) -> ShuffleNetV2:
@@ -346,7 +346,7 @@ def shufflenet_v2_x1_5(
346346
return _shufflenetv2(weights, progress, [4, 8, 4], [24, 176, 352, 704, 1024], **kwargs)
347347

348348

349-
@handle_legacy_interface(weights=("pretrained", None))
349+
@handle_legacy_interface(weights=("pretrained", ShuffleNet_V2_X2_0_Weights.IMAGENET1K_V1))
350350
def shufflenet_v2_x2_0(
351351
*, weights: Optional[ShuffleNet_V2_X2_0_Weights] = None, progress: bool = True, **kwargs: Any
352352
) -> ShuffleNetV2:

0 commit comments

Comments
 (0)