@Fabio Grätz - I think the problem is torch.elastic uses

pickle

instead of

cloudpickle

this works

from flytekit import task
import cloudpickle

@task(cache=True, cache_version="1.0")
def foo(i: int) -> int:
    return i

foo(i=10)
p = cloudpickle.dumps(foo)
f = cloudpickle.loads(p)
f(i=10)