torch.Size error

[smcch]

GaNDLF Version
Version: 0.0.16.

Desktop (please complete the following information):

• OS: Ubuntu (WSL)

How did you install GaNDLF
I followed the installation guide

Dataset description
MRI scans, 4 modalities, 240 x 240 x 155

Describe your question/problem
I got this error after run training:

UserWarning: Using a target size (torch.Size([1, 1])) that is different to the input size (torch.Size([1, 2, 64, 64, 64])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.
loss = F.mse_loss(prediction, target, reduction=reduction)
Looping over training data: 0%| | 0/125 [00:03<?, ?it/s]
ERROR: shape '[1, 1]' is invalid for input of size 262144

My config file:

# affix version
version:
  {
    minimum: 0.0.16,
    maximum: 0.0.16 # this should NOT be made a variable, but should be tested after every tag is created
  }
## Choose the model parameters here
model:
  {
    dimension: 3, # the dimension of the model and dataset: defines dimensionality of computations
    base_filters: 30, # Set base filters: number of filters present in the initial module of the U-Net convolution; for IncU-Net, keep this divisible by 4
    architecture: resunet, # options: unet, resunet, deep_resunet, deep_unet, light_resunet, light_unet, fcn, uinc, vgg, densenet
    norm_type: batch, # options: batch, instance, or none (only for VGG); used for all networks
    final_layer: softmax, # can be either sigmoid, softmax or none (none == regression/logits)
    class_list: ['0','1'], # Set the list of labels the model should train on and predict
    amp: False, # Set if you want to use Automatic Mixed Precision for your operations or not - options: True, False
    num_channels: 4, # set the input channels - useful when reading RGB or images that have vectored pixel types from the CSV
    print_summary: True, # prints the summary of the model before training; defaults to True


## metrics to evaluate the validation performance
metrics:
    - accuracy # classification
    - classification_accuracy # classification
    - per_label_accuracy # used for classification
    - f1 # classification/segmentation
    - precision # classification/segmentation ## more details https://torchmetrics.readthedocs.io/en/latest/references/modules.html#id3
    - recall # classification/segmentation ## more details 
inference_mechanism: {
  grid_aggregator_overlap: crop, # this option provides the option to strategize the grid aggregation output; should be either 'crop' or 'average' - https://torchio.readthedocs.io/patches/patch_inference.html#grid-aggregator
  patch_overlap: 0, # amount of overlap of patches during inference, defaults to 0; see https://torchio.readthedocs.io/patches/patch_inference.html#gridsampler
}
# this is to enable or disable lazy loading - setting to true reads all data once during data loading, resulting in improvements
# in I/O at the expense of memory consumption
in_memory: False
# this will save the generated masks for validation and testing data for qualitative analysis
save_output: False
# this will save the patches used during training for qualitative analysis
save_training: False
# Set the Modality : rad for radiology, path for histopathology
modality: rad
## Patch size during training - 2D patch for breast images since third dimension is not patched
patch_size: [64,64,64]
# uniform: UniformSampler or label: LabelSampler
patch_sampler: uniform
enable_padding: False
# Number of epochs
num_epochs: 10
# Set the patience - measured in number of epochs after which, if the performance metric does not improve, exit the training loop - defaults to the number of epochs
patience: 2
# Set the batch size
batch_size: 1
# gradient clip : norm, value, agc
clip_mode: norm
# clip_gradient value
clip_grad: 0.1
## Set the initial learning rate
learning_rate: 0.001
# Learning rate scheduler - options:"triangle", "triangle_modified", "exp", "step", "reduce-on-plateau", "cosineannealing", "triangular", "triangular2", "exp_range"
# triangle/triangle_modified use LambdaLR but triangular/triangular2/exp_range uses CyclicLR
scheduler:
  {
    type: triangle,
    min_lr: 0.00001,
    max_lr: 1,
  }
# Set which loss function you want to use - options : 'dc' - for dice only, 'dcce' - for sum of dice and CE and you can guess the next (only lower-case please)
# options: dc (dice only), dc_log (-log of dice), ce (), dcce (sum of dice and ce), mse () ...
# mse is the MSE defined by torch and can define a variable 'reduction'; see https://pytorch.org/docs/stable/generated/torch.nn.MSELoss.html#torch.nn.MSELoss
# use mse_torch for regression/classification problems and dice for segmentation
loss_function: mse
# this parameter weights the loss to handle imbalanced losses better
weighted_loss: True

optimizer: adam

nested_training:
  {
    testing: 5, # this controls the testing data splits for final model evaluation; use '1' if this is to be disabled
    validation: 5 # this controls the validation data splits for model training
  }

data_augmentation:
  {
    'flip':{
      'axis': [0,1,2] # one or more axes can be put here. if this isn't defined, all axes are considered
    },
    'rotate_180', # explicitly rotate image by 180; if 'axis' isn't defined, default is [1,2,3]

q_max_length: 40
# this determines the number of patches to extract from each volume. A small number of patches ensures a large variability in the queue, but training will be slower
q_samples_per_volume: 5
# this determines the number subprocesses to use for data loading; '0' means main process is used
q_num_workers: 2 # scale this according to available CPU resources
# used for debugging
q_verbose: False