I haven't managed to successfully export my custom ViT model yet, but I've not had an issue accessing the export methods in torch 2.3 within the nvcr.io/nvidia/

I haven't managed to successfully export my custom ViT model yet, but I've not had an issue accessing the export methods in torch 2.3 within the nvcr.io/nvidia/pytorch:24.02-py3 container. I may have some more time to debug my trace tonight (i.e. remove conditionals from model + make sure everything is on CPU) and will update if I have any new insights. ``` from torch.export import export ... example_args = (dummy_input) exported_program = export(model, args=example_args) ``` Links: - torch.export docs: https://pytorch.org/docs/stable/export.html#serialization https://pytorch.org/docs/stable/export.html#serialization - Using 24.02 container: https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-02.html https://docs.nvidia.com/deeplearning/frameworks/pytorch-rele...