fix: address review comments (iteration #1)

Author: aviruthen

sagemaker-core/pyproject.toml

@@ -66,6 +66,7 @@ test = [
     "pytest>=8.0.0, <9.0.0",
     "pytest-cov>=4.0.0",
     "pytest-xdist>=3.0.0",
+    "sagemaker-core[torch]",
 ]
 
 [project.urls]

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

@@ -445,13 +445,14 @@ def __init__(self, content_type="tensor/pt"):
         super(TorchTensorSerializer, self).__init__(content_type=content_type)
         try:
             from torch import Tensor
+
+            self.torch_tensor = Tensor
         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()
 
     def serialize(self, data):

sagemaker-core/tests/unit/test_serializer_implementations.py

@@ -11,7 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """Unit tests for sagemaker.core.serializers.implementations module."""
-from __future__ import absolute_import
+from __future__ import annotations
 
 import pytest
 from unittest.mock import Mock, patch
@@ -162,3 +162,25 @@ def test_numpy_serializer_import(self):
     def test_record_serializer_deprecated(self):
         """Test that numpy_to_record_serializer is available as deprecated."""
         assert hasattr(implementations, "numpy_to_record_serializer")
+
+
+class TestTorchSerializerWithOptionalDependency:
+    """Test torch serializer/deserializer with optional torch dependency."""
+
+    def test_torch_tensor_serializer_instantiation(self):
+        """Test that TorchTensorSerializer can be instantiated when torch is available."""
+        torch = pytest.importorskip("torch")
+        from sagemaker.core.serializers.base import TorchTensorSerializer
+
+        serializer = TorchTensorSerializer()
+        assert serializer is not None
+        assert serializer.content_type == "tensor/pt"
+
+    def test_torch_tensor_deserializer_instantiation(self):
+        """Test that TorchTensorDeserializer can be instantiated when torch is available."""
+        torch = pytest.importorskip("torch")
+        from sagemaker.core.deserializers.base import TorchTensorDeserializer
+
+        deserializer = TorchTensorDeserializer()
+        assert deserializer is not None
+        assert deserializer.accept == "tensor/pt"

sagemaker-core/tox.ini

@@ -6,7 +6,6 @@
 [tox]
 isolated_build = true
 envlist = black-format,flake8,pylint,docstyle,sphinx,doc8,twine,py39,py310,py311,py312
-
 skip_missing_interpreters = False
 
 [flake8]