Got error with 3D CNN

[zhonhel]

I got error, when using 3D U-Net, please clone this repo to check: https://github.com/wolny/pytorch-3dunet

I use pytorch 1.9.0

Look at this repo
First, change the file_paths item in ./resources/train_config_dice.yaml, with "{path to}/resources/sample_ovule.h5", both for train and val ( line 82 and line 139 )
Then, run use python {path to}/pytorch3dunet/train.py --config {path to}/resources/train_config_dice.yaml

I inserted codes:

    from torchsummary import summary
    summary(self.model, (input.shape[1], input.shape[2], input.shape[3], input.shape[4]))
    exit()

before

    output = self.model(input)

in UNet3DTrainer._forward_pass(self, input, target, weight=None) this function