Issue of input data structure with 3 orthogonal bval

[hylin-work]

I would like to train a de-noise model using our DWI dataset with b0, b500 and b1000 in axial, sagittal and coronal axes. The loaded data size is (256, 256, 38, 3). I am not sure what parameters I have to change to make it work according this data structure. The changes I made from sample hardi_150.json are dataroot, val_volme_idx, val_slice_idx as attached. However, the stage_1 is running forever because of my new settings.

Any suggestions to fix this error is strongly welcome.

config.json
"datasets": {
"train": {
"name": "head",
"dataroot": "data/DWI_HEAD/HEAD.nii",
"valid_mask": [0,3], //[10,160]
"phase": "train",
"padding": 3,
"val_volume_idx": 2,// the volume to visualize for validation
"val_slice_idx": 20, // the slice to visualize for validation
"batch_size": 32,
"in_channel": 1,
"num_workers": 0,
"use_shuffle": true
},
"val": {
"name": "head",
"dataroot": "data/DWI_HEAD/HEAD.nii",
"valid_mask": [0,3], //[10,160]
"phase": "val",
"padding": 3
"val_volume_idx": 2, // the volume to visualize for validation
"val_slice_idx": 20, // the slice to visualize for validation
"batch_size": 1,
"in_channel": 1,
"num_workers": 0
}

"noise_model": {
"resume_state": null,
"unet": {
"in_channel": 2,
"out_channel": 1,
"inner_channel": 32,
"norm_groups": 32,
"channel_multiplier": [
1,
2,
4,
8,
8
],
"attn_res": [
16
],
"res_blocks": 2,
"dropout": 0.0,
"version": "v1"
},

"stage2_file": null

HEAD.bval
0 500 1000

HEAD.bvec
1 0 0
0 1 0
0 0 1

[tiangexiang]

Are there any errors/outputs/logs when running stage_1? I may need more information to help identify the problem. The config set up looks good to me