nan in the training loss

[Hanyi11]

• Python version: 3.9
• Operating System: linux

Description

When training the membrain-seg model, it can happen that sometimes the training loss and training surface dice become nan.

git clone https://github.com/teamtomo/membrain-seg.git
%cd membrain-seg
!pip install .
%cd membrain-seg
!membrain
!membrain train
conda activate memBrain
membrain train --data-dir /home/icb/hanyi.zhang/534_withnewSyntheticAndCollaborators
membrain train_advanced --num-workers 16 --data-dir /home/icb/hanyi.zhang/534_withnewSyntheticAndCollaborators

Epoch 21: 100%
EPOCH Training loss nan
EPOCH Training acc 0.9495545029640198
EPOCH Training surface dice nan