we should add support for easy exportation of models, maybe even export with torch.compile or similar