Hi, I'm trying to register a flow and getting pytorch error. Any input is greatly appreciated. @high-park-82026 @freezing-airport-6809@tall-lock-23197 (XADM-1068/Code$ pyflyte register --image openaimp/flyte:1 Traceback (most recent call last): File "/home/xaimpl/.local/bin/pyflyte", line 5, in <module> from flytekit.clis.sdk_in_container.pyflyte import main File "/home/xaimpl/.local/lib/python3.10/site-packages/flytekit/__init__.py", line 185, in <module> from flytekit.extras import pytorch File "/home/xaimpl/.local/lib/python3.10/site-packages/flytekit/extras/pytorch/__init__.py", line 26, in <module> from .checkpoint import PyTorchCheckpoint, PyTorchCheckpointTransformer File "/home/xaimpl/.local/lib/python3.10/site-packages/flytekit/extras/pytorch/checkpoint.py", line 25, in <module> class PyTorchCheckpoint: File "/home/xaimpl/.local/lib/python3.10/site-packages/flytekit/extras/pytorch/checkpoint.py", line 30, in PyTorchCheckpoint module: Optional[torch.nn.Module] = None

AttributeError: module 'torch' has no attribute 'nn'

@nutritious-flower-32708, have you created a torch.py file?

do you not have pytorch installed?

@tall-lock-23197 no, I've not created a torch.py file. Can you guide me or share a sample file? Also, where that file needs to go?

Is

torch.nn

working in your Python interpreter?

yes

It works for me, @nutritious-flower-32708. Also I noticed you’re not giving a package or a module to pyflyte register command. It has to be of this format:

pyflyte register --image <http://ghcr.io/flyteorg/flytecookbook:core-latest|ghcr.io/flyteorg/flytecookbook:core-latest> pytorch_types.py

.