with a bit of googling, I am guessing its trying to link to this section here? https://www.union.ai/docs/flyte/tutorials/model-training/mnist-classifier/#pytorch-dockerfile-for-deployment

I'm slightly confused, do I need to manage dockerfiles for all of my model training tasks that require gpu dependencies?

I see in this section of the ImageSpec docs that you can use nvidia docker image as a base: https://www.union.ai/docs/flyte/user-guide/core-concepts/image-spec/#install-cuda-in-the-image

But this Github issue to support GPU dependencies via ImageSpec is still open: https://github.com/flyteorg/flyte/issues/3771