Repository: mlcommons/GaNDLF

Issue:
Bug in regression causing different output size to get broadcasted
**Describe the bug**
The regression code is giving a weird error.

**To Reproduce**
Steps to reproduce the behavior:
1. Run the command:
```powershell
conda activate ./venv
pytest --device cuda
```
2. See this error:
```powershell
##vso[task.logissue type=warning;]Using a target size (torch.Size([1])) that is different to the input size (torch.Size([2])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.
```

**Expected behavior**
This error should not be present, because it is causing issues with regression.

**Screenshots**
N.A.

**GaNDLF Version**
0.0.10-dev

**Desktop (please complete the following information):**
N.A.

**Additional context**
N.A.



Discussion:

Resolving patch:
diff --git a/GANDLF/utils.py b/GANDLF/utils.py
--- a/GANDLF/utils.py
+++ b/GANDLF/utils.py
@@ -10,6 +10,7 @@
 import torch
 import torch.nn as nn
 import torchio
+from GANDLF.models.modelBase import get_final_layer
 
 
 def resample_image(
@@ -155,7 +156,7 @@ def reverse_one_hot(predmask_array, class_list):
     final_mask = 0
     special_cases_to_check = ["||"]
     special_case_detected = False
-    max = 0
+    max_current = 0
 
     for _class in class_list:
         for case in special_cases_to_check:
@@ -166,8 +167,8 @@ def reverse_one_hot(predmask_array, class_list):
                         case
                     )  # if present, then split the sub-class
                     for i in class_split:  # find the max for computation later on
-                        if int(i) > max:
-                            max = int(i)
+                        if int(i) > max_current:
+                            max_current = int(i)
 
     if special_case_detected:
         start_idx = 0
@@ -183,7 +184,7 @@ def reverse_one_hot(predmask_array, class_list):
                 predmask_array[0, :, :, :], dtype=int
             )  # predmask_array[i,:,:,:].long()
             # temp_sum = torch.sum(output)
-        # output_2 = (max - torch.sum(output)) % max
+        # output_2 = (max_current - torch.sum(output)) % max_current
         # test_2 = 1
     else:
         for idx, _class in enumerate(class_list):
@@ -335,7 +336,7 @@ def populate_header_in_parameters(parameters, headers):
     if len(headers["predictionHeaders"]) > 0:
         parameters["model"]["num_classes"] = len(headers["predictionHeaders"])
     is_regression, _, _ = find_problem_type(
-        parameters["headers"], parameters["model"]["final_layer"]
+        parameters["headers"], get_final_layer(parameters["model"]["final_layer"])
     )
 
     # if the problem type is classification/segmentation, ensure the number of classes are picked from the configuration
diff --git a/GANDLF/version.py b/GANDLF/version.py
--- a/GANDLF/version.py
+++ b/GANDLF/version.py
@@ -1,3 +1,3 @@
 #!/usr/bin/env python3
 # -*- coding: UTF-8 -*-
-__version__ = "0.0.10-dev"
+__version__ = "0.0.10"


Test patch:
diff --git a/testing/config_classification.yaml b/testing/config_classification.yaml
--- a/testing/config_classification.yaml
+++ b/testing/config_classification.yaml
@@ -17,7 +17,7 @@ model:
   - 1
   - 2
   dimension: 2
-  final_layer: None
+  final_layer: softmax
 nested_training:
   testing: -2
   validation: -2