Fix PyTorch transformer

flytekit/extras/pytorch/native.py

@@ -1,5 +1,5 @@
 import pathlib
-from typing import Generic, Type, TypeVar
+from typing import Type, TypeVar
 
 import torch
 
@@ -12,7 +12,7 @@
 T = TypeVar("T")
 
 
-class PyTorchTypeTransformer(TypeTransformer, Generic[T]):
+class PyTorchTypeTransformer(TypeTransformer[T]):
     def get_literal_type(self, t: Type[T]) -> LiteralType:
         return LiteralType(
             blob=_core_types.BlobType(
@@ -63,30 +63,40 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type:
         # load pytorch tensor/module from a file
         return torch.load(local_path, map_location=map_location)
 
-    def guess_python_type(self, literal_type: LiteralType) -> Type[T]:
+
+class PyTorchTensorTransformer(PyTorchTypeTransformer[torch.Tensor]):
+    PYTORCH_FORMAT = "PyTorchTensor"
+
+    def __init__(self):
+        super().__init__(name="PyTorch Tensor", t=torch.Tensor)
+
+    def guess_python_type(self, literal_type: LiteralType) -> Type[torch.Tensor]:
         if (
             literal_type.blob is not None
             and literal_type.blob.dimensionality == _core_types.BlobType.BlobDimensionality.SINGLE
             and literal_type.blob.format == self.PYTORCH_FORMAT
         ):
-            return T
+            return torch.Tensor
 
         raise ValueError(f"Transformer {self} cannot reverse {literal_type}")
 
 
-class PyTorchTensorTransformer(PyTorchTypeTransformer[torch.Tensor]):
-    PYTORCH_FORMAT = "PyTorchTensor"
-
-    def __init__(self):
-        super().__init__(name="PyTorch Tensor", t=torch.Tensor)
-
-
 class PyTorchModuleTransformer(PyTorchTypeTransformer[torch.nn.Module]):
     PYTORCH_FORMAT = "PyTorchModule"
 
     def __init__(self):
         super().__init__(name="PyTorch Module", t=torch.nn.Module)
 
+    def guess_python_type(self, literal_type: LiteralType) -> Type[torch.nn.Module]:
+        if (
+            literal_type.blob is not None
+            and literal_type.blob.dimensionality == _core_types.BlobType.BlobDimensionality.SINGLE
+            and literal_type.blob.format == self.PYTORCH_FORMAT
+        ):
+            return torch.nn.Module
+
+        raise ValueError(f"Transformer {self} cannot reverse {literal_type}")
+
 
 TypeEngine.register(PyTorchTensorTransformer())
 TypeEngine.register(PyTorchModuleTransformer())

tests/flytekit/unit/extras/pytorch/test_transformations.py

@@ -40,6 +40,7 @@ def test_get_literal_type(transformer, python_type, format):
     tf = transformer
     lt = tf.get_literal_type(python_type)
     assert lt == LiteralType(blob=BlobType(format=format, dimensionality=BlobType.BlobDimensionality.SINGLE))
+    assert tf.guess_python_type(lt) == python_type
 
 
 @pytest.mark.parametrize(