[BUG] Can't get torch.nn.Sequential from FlyteRemote: "Type ~T not supported currently in Flytekit. Please register a new transformer"

[cosmicBboy]

Describe the bug

Trying to get a FlyteExecution output with execution.outputs["o0"] where the output type is torch.nn.Sequential (a subclass of torch.nn.Module`) raises the following error:

Traceback (most recent call last):
  File "/Users/nielsbantilan/git/flyte-demos/flyte_demo/test.py", line 23, in <module>
    print(execution.outputs["o0"])
  File "/Users/nielsbantilan/miniconda3/envs/flyte-demos/lib/python3.9/site-packages/flytekit/core/type_engine.py", line 1584, in __getitem__
    return self.get(key)
  File "/Users/nielsbantilan/miniconda3/envs/flyte-demos/lib/python3.9/site-packages/flytekit/core/type_engine.py", line 1613, in get
    val = TypeEngine.to_python_value(self._ctx or FlyteContext.current_context(), self._literals[attr], as_type)
  File "/Users/nielsbantilan/miniconda3/envs/flyte-demos/lib/python3.9/site-packages/flytekit/core/type_engine.py", line 763, in to_python_value
    transformer = cls.get_transformer(expected_python_type)
  File "/Users/nielsbantilan/miniconda3/envs/flyte-demos/lib/python3.9/site-packages/flytekit/core/type_engine.py", line 702, in get_transformer
    raise ValueError(f"Type {python_type} not supported currently in Flytekit. Please register a new transformer")
ValueError: Type ~T not supported currently in Flytekit. Please register a new transformer

Need to set this explicitly with execution.outputs.get("o0", as_type=torch.nn.Sequential).

Expected behavior

The flyte type engine should be able to get supported outputs without having to specify as_type.

Additional context to reproduce

To reproduce:

• git clone https://github.com/flyteorg/flyte-demos
• pyflyte run the workflows/model_training.py script: https://github.com/flyteorg/flyte-demos/blob/main/flyte_demo/workflows/model_training.py
• run the following code snippet in the flyte_demo project directory

from flytekit.configuration import Config
from flytekit.remote import FlyteRemote
from workflows import model_training

remote = FlyteRemote(
    config=Config.auto(),
    default_project="flytesnacks",
    default_domain="development",
)
execution = remote.execute_local_workflow(
    model_training.training_workflow,
    inputs={
        "n_epochs": 30,
        "hyperparameters": model_training.Hyperparameters(
            in_dim=4, hidden_dim=100, out_dim=3, learning_rate=0.01
        ),
    }
)

print(remote.generate_console_url(execution))
execution = remote.wait(execution)
print(execution.outputs["o0"])

Screenshots

No response

Are you sure this issue hasn't been raised already?

• Yes

Have you read the Code of Conduct?

• Yes