Unable to plot network architecture with hiddenlayer

[Haneskixy]

Hello !
I want to plot the networks with the help of hiddenlayer. I've tried on both nnunetv1 and nnunetv2 but neither of them worked. I've got this error: Unable to plot network architecture:module 'torch.onnx' has no attribute '_optimize_trace'.
My CUDA version is 11.0 and these are the packages I've installed in my virtual environment:
asttokens 2.4.1
batchgenerators 0.25
certifi 2024.6.2
charset-normalizer 3.3.2
contourpy 1.2.1
cycler 0.12.1
decorator 5.1.1
dicom2nifti 2.4.11
exceptiongroup 1.2.1
executing 2.0.1
filelock 3.14.0
fonttools 4.53.0
fsspec 2024.6.0
future 1.0.0
graphviz 0.20.1
hiddenlayer 0.3
idna 3.7
imageio 2.34.1
importlib_resources 6.4.0
ipython 8.18.1
jedi 0.19.1
Jinja2 3.1.4
joblib 1.4.2
kiwisolver 1.4.5
lazy_loader 0.4
linecache2 1.0.0
MarkupSafe 2.1.5
matplotlib 3.9.0
matplotlib-inline 0.1.7
MedPy 0.5.1
mpmath 1.3.0
networkx 3.2.1
nibabel 5.2.1
nnunet 1.7.1
numpy 1.26.4
nvidia-cublas-cu12 12.1.3.1
nvidia-cuda-cupti-cu12 12.1.105
nvidia-cuda-nvrtc-cu12 12.1.105
nvidia-cuda-runtime-cu12 12.1.105
nvidia-cudnn-cu12 8.9.2.26
nvidia-cufft-cu12 11.0.2.54
nvidia-curand-cu12 10.3.2.106
nvidia-cusolver-cu12 11.4.5.107
nvidia-cusparse-cu12 12.1.0.106
nvidia-nccl-cu12 2.20.5
nvidia-nvjitlink-cu12 12.5.40
nvidia-nvtx-cu12 12.1.105
packaging 24.0
pandas 2.2.2
parso 0.8.4
pexpect 4.9.0
pillow 10.3.0
pip 24.0
prompt_toolkit 3.0.45
ptyprocess 0.7.0
pure-eval 0.2.2
pydicom 2.4.4
Pygments 2.18.0
pyparsing 3.1.2
python-dateutil 2.9.0.post0
python-gdcm 3.0.24.1
pytz 2024.1
requests 2.32.3
scikit-image 0.22.0
scikit-learn 1.5.0
scipy 1.13.1
setuptools 69.5.1
SimpleITK 2.3.1
six 1.16.0
stack-data 0.6.3
sympy 1.12.1
threadpoolctl 3.5.0
tifffile 2024.5.22
torch 2.3.0
torchaudio 0.7.2
torchvision 0.8.2+cu110
tqdm 4.66.4
traceback2 1.4.0
traitlets 5.14.3
triton 2.3.0
typing_extensions 4.12.1
tzdata 2024.1
unittest2 1.1.0
urllib3 2.2.1
wcwidth 0.2.13
wheel 0.43.0
zipp 3.19.1
Do you know how to fix this problem please ? Thank you in advance !

[kovecsroli]

What did you try to run?