Crash when using ONNX file with NNI plugin

Hi, I’m trying to import an ONNX file that I have just exported from PyTorch. I have the NNI plugin enabled, and I understand that I should just be able to drag the ONNX file into the project and it should be recognised. Instead UE5.2 crashes with the message:

Assertion failed: false [File:D:\build++UE5\Sync\Engine\Plugins\Experimental\NNI\Source\Private\ThirdPartyHelperAndDLLLoader\Private\ORTExceptionHandler.cpp] [Line: 16] ONNXRuntime threw an exception with code 6, e.what(): "Exception during initialization: D:\build++UE5\Sync\Engine\Plugins\Experimental\NNI\Source\ThirdParty\ORT_4_1\Private\Windows\core\providers\dml\DmlExecutionProvider\src\GraphKernelHelper.cpp(78)\UnrealEditor-ONNXRuntime.dll!00007FFBF2FB5E89: (caller: 00007FFBF2FF760C) Exception(1) tid(2a60) 80070057 The parameter is incorrect. ".

Is there something special I need to do when exporting from PyTorch in order for it to be reconised properly in UE5?

Thanks

An update on this: it appears that if the network is too big, PyTorch will export as multiple files for the weights and biases and a small onnx file. I believe this may be something that NNI can’t handle. I significantly reduced my number of hidden neurons and the export went to a single file and UE5 loaded it successfully.