Skip to content

Commit 3fc162e

Browse files
authored
fix: make PyTorch empty framework_version warning include the latest PyTorch version (#1209)
1 parent c9ee6c4 commit 3fc162e

File tree

2 files changed

+4
-4
lines changed

2 files changed

+4
-4
lines changed

src/sagemaker/pytorch/estimator.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright 2018-2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
1+
# Copyright 2018-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
22
#
33
# Licensed under the Apache License, Version 2.0 (the "License"). You
44
# may not use this file except in compliance with the License. A copy of
@@ -104,7 +104,7 @@ def __init__(
104104
:class:`~sagemaker.estimator.EstimatorBase`.
105105
"""
106106
if framework_version is None:
107-
logger.warning(empty_framework_version_warning(PYTORCH_VERSION, PYTORCH_VERSION))
107+
logger.warning(empty_framework_version_warning(PYTORCH_VERSION, self.LATEST_VERSION))
108108
self.framework_version = framework_version or PYTORCH_VERSION
109109

110110
if "enable_sagemaker_metrics" not in kwargs:

tests/unit/test_pytorch.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright 2017-2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
1+
# Copyright 2017-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
22
#
33
# Licensed under the Apache License, Version 2.0 (the "License"). You
44
# may not use this file except in compliance with the License. A copy of
@@ -529,7 +529,7 @@ def test_empty_framework_version(warning, sagemaker_session):
529529
)
530530

531531
assert estimator.framework_version == defaults.PYTORCH_VERSION
532-
warning.assert_called_with(defaults.PYTORCH_VERSION, defaults.PYTORCH_VERSION)
532+
warning.assert_called_with(defaults.PYTORCH_VERSION, estimator.LATEST_VERSION)
533533

534534

535535
def test_pt_enable_sm_metrics(sagemaker_session):

0 commit comments

Comments
 (0)