fix: address review comments (iteration #1)

Author: aviruthen

sagemaker-core/pyproject.toml

@@ -60,7 +60,7 @@ torch = [
     "torch>=1.9.0",
 ]
 all = [
-    "torch>=1.9.0",
+    "sagemaker-core[torch]",
 ]
 test = [
     "pytest>=8.0.0, <9.0.0",

sagemaker-core/src/sagemaker/core/deserializers/base.py

@@ -365,11 +365,11 @@ def __init__(self, accept="tensor/pt"):
             from torch import from_numpy
 
             self.convert_npy_to_tensor = from_numpy
-        except ImportError:
+        except ImportError as e:
             raise ImportError(
                 "Unable to import torch. Please install torch to use TorchTensorDeserializer: "
                 "pip install 'sagemaker-core[torch]'"
-            )
+            ) from e
 
     def deserialize(self, stream, content_type="tensor/pt"):
         """Deserialize streamed data to TorchTensor

sagemaker-core/src/sagemaker/core/serializers/base.py

@@ -445,11 +445,11 @@ def __init__(self, content_type="tensor/pt"):
         super(TorchTensorSerializer, self).__init__(content_type=content_type)
         try:
             from torch import Tensor
-        except ImportError:
+        except ImportError as e:
             raise ImportError(
                 "Unable to import torch. Please install torch to use TorchTensorSerializer: "
                 "pip install 'sagemaker-core[torch]'"
-            )
+            ) from e
 
         self.torch_tensor = Tensor
         self.numpy_serializer = NumpySerializer()

sagemaker-core/tests/unit/test_torch_optional_dependency.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 import sys
-from unittest.mock import patch, MagicMock
+from unittest.mock import patch
 
 import numpy as np
 import pytest
@@ -36,7 +36,7 @@ def test_torch_tensor_serializer_raises_import_error_when_torch_missing():
 
 
 def test_torch_tensor_deserializer_raises_import_error_when_torch_missing():
-    """Verify TorchTensorDeserializer raises ImportError with helpful message when torch is missing."""
+    """Verify TorchTensorDeserializer raises ImportError when torch is missing."""
     import importlib
     import sagemaker.core.deserializers.base as base_module
 
@@ -51,6 +51,7 @@ def test_torch_tensor_deserializer_raises_import_error_when_torch_missing():
 
 def test_torch_tensor_serializer_works_when_torch_installed():
     """Verify TorchTensorSerializer can be instantiated when torch is available."""
+    pytest.importorskip("torch")
     from sagemaker.core.serializers.base import TorchTensorSerializer
 
     serializer = TorchTensorSerializer()
@@ -60,6 +61,7 @@ def test_torch_tensor_serializer_works_when_torch_installed():
 
 def test_torch_tensor_deserializer_works_when_torch_installed():
     """Verify TorchTensorDeserializer can be instantiated when torch is available."""
+    pytest.importorskip("torch")
     from sagemaker.core.deserializers.base import TorchTensorDeserializer
 
     deserializer = TorchTensorDeserializer()