We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Perform penalty calculation after all sanity checks are completed
**Is your feature request related to a problem? Please describe.**
The penalty calculation takes a long time, and there are sanity checks that happen after this, which can be a pain.

**Describe the solution you'd like**
It would be great to have these checks before the penalty calculation for quality-of-life improvements.

**Describe alternatives you've considered**
N.A.

**Additional context**
From Evan C.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `GANDLF/compute/generic.py`
Content:
```
1 from GANDLF.models import get_model
2 from GANDLF.schedulers import get_scheduler
3 from GANDLF.optimizers import get_optimizer
4 from GANDLF.data import (
5     get_train_loader,
6     get_validation_loader,
7 )
8 from GANDLF.utils import (
9     populate_header_in_parameters,
10     parseTrainingCSV,
11     send_model_to_device,
12     get_class_imbalance_weights,
13 )
14 
15 
16 def create_pytorch_objects(parameters, train_csv=None, val_csv=None, device="cpu"):
17     """
18     This function creates all the PyTorch objects needed for training.
19 
20     Args:
21         parameters (dict): The parameters dictionary.
22         train_csv (str): The path to the training CSV file.
23         val_csv (str): The path to the validation CSV file.
24         device (str): The device to perform computations on.
25 
26     Returns:
27         model (torch.nn.Module): The model to use for training.
28         optimizer (Optimizer): The optimizer to use for training.
29         train_loader (torch.utils.data.DataLoader): The training data loader.
30         val_loader (torch.utils.data.DataLoader): The validation data loader.
31         scheduler (object): The scheduler to use for training.
32         parameters (dict): The updated parameters dictionary.
33     """
34     # initialize train and val loaders
35     train_loader, val_loader = None, None
36     headers_to_populate_train, headers_to_populate_val = None, None
37 
38     if train_csv is not None:
39         # populate the data frames
40         parameters["training_data"], headers_to_populate_train = parseTrainingCSV(
41             train_csv, train=True
42         )
43         parameters = populate_header_in_parameters(
44             parameters, headers_to_populate_train
45         )
46         # get the train loader
47         train_loader = get_train_loader(parameters)
48         parameters["training_samples_size"] = len(train_loader)
49 
50         # Calculate the weights here
51         (
52             parameters["weights"],
53             parameters["class_weights"],
54         ) = get_class_imbalance_weights(parameters["training_data"], parameters)
55 
56     if val_csv is not None:
57         parameters["validation_data"], headers_to_populate_val = parseTrainingCSV(
58             val_csv, train=False
59         )
60         if headers_to_populate_train is None:
61             parameters = populate_header_in_parameters(
62                 parameters, headers_to_populate_val
63             )
64         # get the validation loader
65         val_loader = get_validation_loader(parameters)
66 
67     # get the model
68     model = get_model(parameters)
69     parameters["model_parameters"] = model.parameters()
70 
71     # get the optimizer
72     optimizer = get_optimizer(parameters)
73     parameters["optimizer_object"] = optimizer
74 
75     # send model to correct device
76     (
77         model,
78         parameters["model"]["amp"],
79         parameters["device"],
80         parameters["device_id"],
81     ) = send_model_to_device(
82         model, amp=parameters["model"]["amp"], device=device, optimizer=optimizer
83     )
84 
85     # only need to create scheduler if training
86     if train_csv is not None:
87         if not ("step_size" in parameters["scheduler"]):
88             parameters["scheduler"]["step_size"] = (
89                 parameters["training_samples_size"] / parameters["learning_rate"]
90             )
91 
92         scheduler = get_scheduler(parameters)
93     else:
94         scheduler = None
95 
96     # these keys contain generators, and are not needed beyond this point in params
97     generator_keys_to_remove = ["optimizer_object", "model_parameters"]
98     for key in generator_keys_to_remove:
99         parameters.pop(key, None)
100 
101     return model, optimizer, train_loader, val_loader, scheduler, parameters
102
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/GANDLF/compute/generic.py b/GANDLF/compute/generic.py
--- a/GANDLF/compute/generic.py
+++ b/GANDLF/compute/generic.py
@@ -47,12 +47,6 @@
         train_loader = get_train_loader(parameters)
         parameters["training_samples_size"] = len(train_loader)
 
-        # Calculate the weights here
-        (
-            parameters["weights"],
-            parameters["class_weights"],
-        ) = get_class_imbalance_weights(parameters["training_data"], parameters)
-
     if val_csv is not None:
         parameters["validation_data"], headers_to_populate_val = parseTrainingCSV(
             val_csv, train=False
@@ -90,6 +84,13 @@
             )
 
         scheduler = get_scheduler(parameters)
+
+        # Calculate the weights here
+        (
+            parameters["weights"],
+            parameters["class_weights"],
+        ) = get_class_imbalance_weights(parameters["training_data"], parameters)
+
     else:
         scheduler = None

{"golden_diff": "diff --git a/GANDLF/compute/generic.py b/GANDLF/compute/generic.py\n--- a/GANDLF/compute/generic.py\n+++ b/GANDLF/compute/generic.py\n@@ -47,12 +47,6 @@\n         train_loader = get_train_loader(parameters)\n         parameters[\"training_samples_size\"] = len(train_loader)\n \n-        # Calculate the weights here\n-        (\n-            parameters[\"weights\"],\n-            parameters[\"class_weights\"],\n-        ) = get_class_imbalance_weights(parameters[\"training_data\"], parameters)\n-\n     if val_csv is not None:\n         parameters[\"validation_data\"], headers_to_populate_val = parseTrainingCSV(\n             val_csv, train=False\n@@ -90,6 +84,13 @@\n             )\n \n         scheduler = get_scheduler(parameters)\n+\n+        # Calculate the weights here\n+        (\n+            parameters[\"weights\"],\n+            parameters[\"class_weights\"],\n+        ) = get_class_imbalance_weights(parameters[\"training_data\"], parameters)\n+\n     else:\n         scheduler = None\n", "issue": "Perform penalty calculation after all sanity checks are completed\n**Is your feature request related to a problem? Please describe.**\r\nThe penalty calculation takes a long time, and there are sanity checks that happen after this, which can be a pain.\r\n\r\n**Describe the solution you'd like**\r\nIt would be great to have these checks before the penalty calculation for quality-of-life improvements.\r\n\r\n**Describe alternatives you've considered**\r\nN.A.\r\n\r\n**Additional context**\r\nFrom Evan C.\n", "before_files": [{"content": "from GANDLF.models import get_model\nfrom GANDLF.schedulers import get_scheduler\nfrom GANDLF.optimizers import get_optimizer\nfrom GANDLF.data import (\n    get_train_loader,\n    get_validation_loader,\n)\nfrom GANDLF.utils import (\n    populate_header_in_parameters,\n    parseTrainingCSV,\n    send_model_to_device,\n    get_class_imbalance_weights,\n)\n\n\ndef create_pytorch_objects(parameters, train_csv=None, val_csv=None, device=\"cpu\"):\n    \"\"\"\n    This function creates all the PyTorch objects needed for training.\n\n    Args:\n        parameters (dict): The parameters dictionary.\n        train_csv (str): The path to the training CSV file.\n        val_csv (str): The path to the validation CSV file.\n        device (str): The device to perform computations on.\n\n    Returns:\n        model (torch.nn.Module): The model to use for training.\n        optimizer (Optimizer): The optimizer to use for training.\n        train_loader (torch.utils.data.DataLoader): The training data loader.\n        val_loader (torch.utils.data.DataLoader): The validation data loader.\n        scheduler (object): The scheduler to use for training.\n        parameters (dict): The updated parameters dictionary.\n    \"\"\"\n    # initialize train and val loaders\n    train_loader, val_loader = None, None\n    headers_to_populate_train, headers_to_populate_val = None, None\n\n    if train_csv is not None:\n        # populate the data frames\n        parameters[\"training_data\"], headers_to_populate_train = parseTrainingCSV(\n            train_csv, train=True\n        )\n        parameters = populate_header_in_parameters(\n            parameters, headers_to_populate_train\n        )\n        # get the train loader\n        train_loader = get_train_loader(parameters)\n        parameters[\"training_samples_size\"] = len(train_loader)\n\n        # Calculate the weights here\n        (\n            parameters[\"weights\"],\n            parameters[\"class_weights\"],\n        ) = get_class_imbalance_weights(parameters[\"training_data\"], parameters)\n\n    if val_csv is not None:\n        parameters[\"validation_data\"], headers_to_populate_val = parseTrainingCSV(\n            val_csv, train=False\n        )\n        if headers_to_populate_train is None:\n            parameters = populate_header_in_parameters(\n                parameters, headers_to_populate_val\n            )\n        # get the validation loader\n        val_loader = get_validation_loader(parameters)\n\n    # get the model\n    model = get_model(parameters)\n    parameters[\"model_parameters\"] = model.parameters()\n\n    # get the optimizer\n    optimizer = get_optimizer(parameters)\n    parameters[\"optimizer_object\"] = optimizer\n\n    # send model to correct device\n    (\n        model,\n        parameters[\"model\"][\"amp\"],\n        parameters[\"device\"],\n        parameters[\"device_id\"],\n    ) = send_model_to_device(\n        model, amp=parameters[\"model\"][\"amp\"], device=device, optimizer=optimizer\n    )\n\n    # only need to create scheduler if training\n    if train_csv is not None:\n        if not (\"step_size\" in parameters[\"scheduler\"]):\n            parameters[\"scheduler\"][\"step_size\"] = (\n                parameters[\"training_samples_size\"] / parameters[\"learning_rate\"]\n            )\n\n        scheduler = get_scheduler(parameters)\n    else:\n        scheduler = None\n\n    # these keys contain generators, and are not needed beyond this point in params\n    generator_keys_to_remove = [\"optimizer_object\", \"model_parameters\"]\n    for key in generator_keys_to_remove:\n        parameters.pop(key, None)\n\n    return model, optimizer, train_loader, val_loader, scheduler, parameters\n", "path": "GANDLF/compute/generic.py"}], "after_files": [{"content": "from GANDLF.models import get_model\nfrom GANDLF.schedulers import get_scheduler\nfrom GANDLF.optimizers import get_optimizer\nfrom GANDLF.data import (\n    get_train_loader,\n    get_validation_loader,\n)\nfrom GANDLF.utils import (\n    populate_header_in_parameters,\n    parseTrainingCSV,\n    send_model_to_device,\n    get_class_imbalance_weights,\n)\n\n\ndef create_pytorch_objects(parameters, train_csv=None, val_csv=None, device=\"cpu\"):\n    \"\"\"\n    This function creates all the PyTorch objects needed for training.\n\n    Args:\n        parameters (dict): The parameters dictionary.\n        train_csv (str): The path to the training CSV file.\n        val_csv (str): The path to the validation CSV file.\n        device (str): The device to perform computations on.\n\n    Returns:\n        model (torch.nn.Module): The model to use for training.\n        optimizer (Optimizer): The optimizer to use for training.\n        train_loader (torch.utils.data.DataLoader): The training data loader.\n        val_loader (torch.utils.data.DataLoader): The validation data loader.\n        scheduler (object): The scheduler to use for training.\n        parameters (dict): The updated parameters dictionary.\n    \"\"\"\n    # initialize train and val loaders\n    train_loader, val_loader = None, None\n    headers_to_populate_train, headers_to_populate_val = None, None\n\n    if train_csv is not None:\n        # populate the data frames\n        parameters[\"training_data\"], headers_to_populate_train = parseTrainingCSV(\n            train_csv, train=True\n        )\n        parameters = populate_header_in_parameters(\n            parameters, headers_to_populate_train\n        )\n        # get the train loader\n        train_loader = get_train_loader(parameters)\n        parameters[\"training_samples_size\"] = len(train_loader)\n\n    if val_csv is not None:\n        parameters[\"validation_data\"], headers_to_populate_val = parseTrainingCSV(\n            val_csv, train=False\n        )\n        if headers_to_populate_train is None:\n            parameters = populate_header_in_parameters(\n                parameters, headers_to_populate_val\n            )\n        # get the validation loader\n        val_loader = get_validation_loader(parameters)\n\n    # get the model\n    model = get_model(parameters)\n    parameters[\"model_parameters\"] = model.parameters()\n\n    # get the optimizer\n    optimizer = get_optimizer(parameters)\n    parameters[\"optimizer_object\"] = optimizer\n\n    # send model to correct device\n    (\n        model,\n        parameters[\"model\"][\"amp\"],\n        parameters[\"device\"],\n        parameters[\"device_id\"],\n    ) = send_model_to_device(\n        model, amp=parameters[\"model\"][\"amp\"], device=device, optimizer=optimizer\n    )\n\n    # only need to create scheduler if training\n    if train_csv is not None:\n        if not (\"step_size\" in parameters[\"scheduler\"]):\n            parameters[\"scheduler\"][\"step_size\"] = (\n                parameters[\"training_samples_size\"] / parameters[\"learning_rate\"]\n            )\n\n        scheduler = get_scheduler(parameters)\n\n        # Calculate the weights here\n        (\n            parameters[\"weights\"],\n            parameters[\"class_weights\"],\n        ) = get_class_imbalance_weights(parameters[\"training_data\"], parameters)\n\n    else:\n        scheduler = None\n\n    # these keys contain generators, and are not needed beyond this point in params\n    generator_keys_to_remove = [\"optimizer_object\", \"model_parameters\"]\n    for key in generator_keys_to_remove:\n        parameters.pop(key, None)\n\n    return model, optimizer, train_loader, val_loader, scheduler, parameters\n", "path": "GANDLF/compute/generic.py"}]}