[Training] Cannot generate ONNX graph for training

[srijanie03]

Describe the issue

I am trying to generate a training graph by specifying the model path directly. It works correctly on CUDA (12.4) but when I run the same on Windows 11, I get an error as this:

AttributeError: 'str' object has no attribute 'graph'.
When I print the forward graph I can clearly see the attribute 'graph' in it.

It does not seem to be model specific as this issue occurs for every graph I try by specifying the model path. Doing onnx.load(onnx_model_path) works but not for a model that is bigger than 2GB.

To reproduce:

https://github.com/microsoft/onnxruntime-training-examples/blob/master/on_device_training/desktop/python/mnist.ipynb

#Backward graph artifacts.generate_artifacts( "mnist.onnx", optimizer=artifacts.OptimType.AdamW, loss=artifacts.LossType.CrossEntropyLoss, requires_grad=requires_grad, frozen_params=frozen_params, artifact_directory="mnist", additional_output_names=["output"])

Urgency

Very urgent

ONNX Runtime Installation

Built from Source

ONNX Runtime Version or Commit ID

1.19.0

PyTorch Version

2.4.0

Execution Provider

Default CPU

Execution Provider Library Version

CUDA 12.4