correct mistaie in previous commit (#528)

In TF tests, verify that TF version is supported, not pytorch's

Co-authored-by: Adam Cherti <[email protected]>

Author: adimux

tests/pytorch/test_get_hook.py

@@ -1,9 +1,11 @@
+# Third Party
 import pytest
 import torch
 from numpy import percentile
 from tests.utils import SagemakerSimulator, Timer, is_running_in_codebuild
 from torch.utils.smdebug import get_smdebug_hook
 
+# First Party
 from smdebug.core.utils import FRAMEWORK, is_framework_version_supported
 
 
@@ -25,6 +27,7 @@ def test_did_you_forget_to_update_the_supported_framework_version():
             f" if you are trying to release sagemaker-debugger for the next version of pytorch ({torch.__version__})."
         )
 
+
 @pytest.mark.skipif(
     is_running_in_codebuild() is False,
     reason="Microbenchmarking thresholds have been determined only for ci",

tests/tensorflow/keras/test_keras.py

@@ -147,7 +147,7 @@ def test_did_you_forget_to_update_the_supported_framework_version():
 
     This is just to make you aware of the problem.
     """
-    if not is_framework_version_supported(FRAMEWORK.PYTORCH):
+    if not is_framework_version_supported(FRAMEWORK.TENSORFLOW):
         raise Exception(
             "You are running against an unsupported version of Tensorflow..."
             " Please update `smdebug.tensorflow.utils.SUPPORTED_TF_VERSION_THRESHOLD`"