diff --git a/README.md b/README.md
index 2efab74..c8e7264 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,7 @@
 # Hyperparameter Tuning with SLURM and Optuna
 
 This package integrates Optuna with SLURM to efficiently run hyperparameter optimization jobs on a cluster. It supports customizable trial parameters and allows you to track the results via a CSV file. The SLURM jobs process the parameter configurations and log their results, which are then read by Optuna to calculate the objective loss.
+
 ## Features
 
 - Submit SLURM jobs with trial parameters generated by Optuna.
@@ -27,7 +28,7 @@ pip install git+https://github.com/C4theBomb/hyperparameter-tuner.git
 from hyperparameter_tuner import Loss
 
 class CustomLossFunction(Loss):
-  def calculate(self, row: pd.Series) -> float:
+  def __call__(self, row: pd.Series, params: Dict[str, Any]) -> float:
     return row[0]
 ```
 
@@ -65,13 +66,14 @@ OPTIMIZER=$5
 REWARDS=$((RANDOM % 100))  # Replace with your experiment logic
 
 # Write results to CSV (first column must be TRIAL_ID so that Optuna can identify the run).
-echo "$TRIAL_ID,$REWARDS,$TRAJECTORIES,$LEARNING_RATE,$OPTIMIZER" >> $RESULTS_PATH
+echo "$TRIAL_ID,$STEP_NUM,$REWARDS,$TRAJECTORIES,$LEARNING_RATE,$OPTIMIZER" >> $RESULTS_PATH
 ```
 
 ### 4. Run your script
 ```
 python main.py
 ```
+
 ## License
 
 [AGPL-3.0](https://choosealicense.com/licenses/agpl-3.0/)
diff --git a/hyperparameter_tuner/slurm_handler.py b/hyperparameter_tuner/slurm_handler.py
index 0ed74de..1d6f484 100644
--- a/hyperparameter_tuner/slurm_handler.py
+++ b/hyperparameter_tuner/slurm_handler.py
@@ -28,13 +28,13 @@ def create_objective(
         slurm_script: str - Path to the SLURM script to execute.
         results_path: str - Path to the CSV file where results will be logged.
         loss: Loss - An instance of a Loss class that implements a `calculate` method.
-        trial_param_types: Dict[str, Tuple[str, Tuple, Dict]]: Dictionary mapping parameter names to their types
-            and arguments. Each entry is structured as:
-            - key str: The name of the parameter.
-            - value: Tuple[str, Tuple, Dict]
-                - str - The parameter type ('int', 'float', or 'categorical').
-                - Tuple - Positional arguments for the parameter's sampling method.
-                - Dict - Keyword arguments for the parameter's sampling method.
+        trial_param_types: Dict[str, Tuple[str, Tuple, Dict]]: Dictionary mapping parameter names to their types and arguments.
+            Each entry is structured as:
+                - key str: The name of the parameter.
+                - value: Tuple[str, Tuple, Dict[str, Any]]
+                    - str - The parameter type ('int', 'float', or 'categorical').
+                    - Tuple - Positional arguments for the parameter's sampling method.
+                    - Dict[str, Any] - Keyword arguments for the parameter's sampling method.
     """
 
     def objective(trial: Trial) -> float: