I’m not entirely sure, sorry about that - it’s been some years since I have needed to train segmentation model. The training code originally came from torchvision: https://github.com/pytorch/vision/tree/main/references/segmentation
I had added support for FCN-ResNet18 and made it ONNX exportable. Then Onixaz packaged it up into a self-contained repo. It appears there have been some updates to the original torchvision code in the meantime (including that error/exception you are getting has now been removed)