N foreground 0, N background 52003401,unable to generate class balanced samples. NaN or Inf found in input tensor.

[shanpriya3]

I am training UNETR model in a custom dataset with labels =0,1,2.

When I train my model using python main.py
--in_channels=1
--out_channels=3
--space_x=1.0
--space_y=1.0
--space_z=1.0
--feature_size=32
--batch_size=1
--logdir=unetrmonai_test
--optim_lr=1e-4
--lrschedule=warmup_cosine
--infer_overlap=0.5
--save_checkpoint
--data_dir=/cluster/home/shan/ves/research-contributions/UNETR/BTCV/dataset/us/

I am getting the following issues, "N foreground 0, N background 52003401,unable to generate class balanced samples. NaN or Inf found in input tensor".

I checked the labels, they are integer values (0,1,2). I am using --out_channels=3. Please let me know what is the issue here. Thanks!