#3783 [BUG] Error with `torch.nn.Module` and `flytekit` version >1.5.0. Issue created by constanceferragu Describe the bug Flyte throws a warning and error with

torch.nn.Module

and

flytekit

version >1.5.0. Flyte uses

PickleFile

with type

nn.Module

as default when importing tasks/workflows. When registering, it does not recognize the type as

nn.Module

but as

PickleFile

- throws an error because the types aren't matching. Example of the error:

# a.py
from flytekit import task

@task
def foo() -> int:
    return 2

------------------------------------------------------------------------------

# b.py
from flytekit import task
from torch import nn

@task
def bla() -> nn.Module:
    return nn.Module()

Warning message:

WARNING  {"asctime": "2023-06-15 17:40:51,991", "name": "flytekit", "levelname": "WARNING", "message": "Unsupported Type <class 'torch.nn.modules.module.Module'> found,      
                                    Flyte will default to use PickleFile as the transport. Pickle can only be used to send objects between the exact same version of Python, and we strongly recommend                   
                                    to use python type that flyte support."}

Error message:

Error 0: Code: MismatchingTypes, Node Id: end-node, Description: Variable [model] (type [blob:<format:"PyTorchModule" > ]) doesn't match expected type [blob:<format:"PythonPickle" > metadata:<fields:<key:"python_class_name" value:<string_value:"Module" > > > ].
Error 1: Code: MismatchingTypes, Node Id: n1, Description: Variable [model] (type [blob:<format:"PythonPickle" > metadata:<fields:<key:"python_class_name" value:<string_value:"Module" > > > ]) doesn't match expected type [blob:<format:"PyTorchModule" > ].

Expected behavior Should be able to register. The culprit is most likely `if is_imported("torch")`: in

flytekit.core.type_engine

Additional context to reproduce No response Screenshots No response Are you sure this issue hasn't been raised already? ☑︎ Yes Have you read the Code of Conduct? ☑︎ Yes flyteorg/flyte