Hey I'm trying to run the pytorch distributed mnist example locally: https://docs.flyte.org/en/latest/flytesnacks/examples/kfpytorch_plugin/pytorch_mnist.html The training seems to go fine but getting the following error:

Failed with Unknown Exception <class 'TypeError'> Reason: Error encountered while executing 'pytorch_training_wf':
  Failed to convert outputs of task 'main.mnist_pytorch_job' at position model_state:
  Incorrect file type, expected application/octet-stream, got application/zip

@gentle-tomato-480 thanks for finding this bug! would you mind updating the docs with this fix? https://github.com/flyteorg/flytesnacks/blob/master/examples/kfpytorch_plugin/kfpytorch_plugin/pytorch_mnist.py#L195

@wide-vegetable-51116 https://github.com/flyteorg/flytesnacks/pull/1687 Done