Good morning. I wish to run a task that utilizes PyTorch, and it will run in a different image than the one that runs the workflow. Do I need to have PyTorch installed in the main image to do so? As I'm currently doing, I need it in both because the main image doesn't recognize PyTorch's import. This is the code I have: from flytekit import task, workflow import torch @task(container_image='') # The container image is omitted def pytorch_test(): # Initialize a tensor x = torch.rand(5, 3) print("Original tensor:", x) # Perform a basic operation (addition) y = x + 5 print("After addition:", y) # Return the modified tensor return y @workflow def pytorch_test_workflow(): result = pytorch_test() return result if name == "__main__": print(pytorch_test_workflow()) How do I run PyTorch and other libraries that the main image doesn't have?

Thank you for the quick answer. By creating a custom image I can run the imports only when I ask for that specific image?

import torch
from flytekit import task, workflow

@task # The container image is omitted
def pytorch_test() -> torch.Tensor:
    # Initialize a tensor
    x = torch.rand(5, 3)
    print("Original tensor:", x)

    # Perform a basic operation (addition)
    y = x + 5
    print("After addition:", y)

    # Return the modified tensor
    return y

@workflow
def pytorch_test_workflow() -> torch.Tensor:
    result = pytorch_test()
    return result

#if __name__ == "__main__":
#   print(pytorch_test_workflow())
pyflyte run {your_wf_file_name}.py pytorch_test_workflow

Here is script check to run in locally first

Thank you for the help, I'll try it out now

@Filipe Fonseca hmm you cannot run all task with 1 image

Yes I'll read it, thank you for explaining

One more way is to delay torch import to inside your function