diff --git a/docs/freqai-parameter-table.md b/docs/freqai-parameter-table.md
index 5e60d2a07..95687c7ab 100644
--- a/docs/freqai-parameter-table.md
+++ b/docs/freqai-parameter-table.md
@@ -100,12 +100,12 @@ Mandatory parameters are marked as **Required** and have to be set in one of the
 
 #### trainer_kwargs
 
-|  Parameter | Description |
-|------------|-------------|
-|  |  **Model training parameters within the `freqai.model_training_parameters.model_kwargs` sub dictionary**
-| `max_iters` | The number of training iterations to run. iteration here refers to the number of times we call self.optimizer.step(). used to calculate n_epochs. <br> **Datatype:** int. <br> Default: `100`.
-| `batch_size` | The size of the batches to use during training.. <br> **Datatype:** int. <br> Default: `64`.
-| `max_n_eval_batches` | The maximum number batches to use for evaluation.. <br> **Datatype:** int, optional. <br> Default: `None`.
+| Parameter    | Description |
+|--------------|-------------|
+|              |  **Model training parameters within the `freqai.model_training_parameters.model_kwargs` sub dictionary**
+| `n_epochs`   | The `n_epochs` parameter is a crucial setting in the PyTorch training loop that determines the number of times the entire training dataset will be used to update the model's parameters. An epoch represents one full pass through the entire training dataset. Overrides `n_steps`. Either `n_epochs` or `n_steps` must be set. <br><br> **Datatype:** int. optional. <br> Default: `10`.
+| `n_steps`    | An alternative way of setting `n_epochs` -  the number of training iterations to run. Iteration here refer to the number of times we call `optimizer.step()`. Ignored if `n_epochs` is set. A simplified version of the function: <br><br> n_epochs = n_steps / (n_obs / batch_size) <br><br> The motivation here is that `n_steps` is easier to optimize and keep stable across different n_obs - the number of data points.  <br> <br> **Datatype:** int. optional. <br> Default: `None`.
+| `batch_size` | The size of the batches to use during training. <br><br> **Datatype:** int. <br> Default: `64`.
 
 
 ### Additional parameters
diff --git a/freqtrade/freqai/prediction_models/PyTorchMLPClassifier.py b/freqtrade/freqai/prediction_models/PyTorchMLPClassifier.py
index 71279dba9..9aabdf7ad 100644
--- a/freqtrade/freqai/prediction_models/PyTorchMLPClassifier.py
+++ b/freqtrade/freqai/prediction_models/PyTorchMLPClassifier.py
@@ -26,9 +26,9 @@ class PyTorchMLPClassifier(BasePyTorchClassifier):
             "model_training_parameters" : {
                 "learning_rate": 3e-4,
                 "trainer_kwargs": {
-                    "max_iters": 5000,
+                    "n_steps": 5000,
                     "batch_size": 64,
-                    "max_n_eval_batches": null,
+                    "n_epochs": null,
                 },
                 "model_kwargs": {
                     "hidden_dim": 512,
diff --git a/freqtrade/freqai/prediction_models/PyTorchMLPRegressor.py b/freqtrade/freqai/prediction_models/PyTorchMLPRegressor.py
index 9f4534487..dc8dc4b61 100644
--- a/freqtrade/freqai/prediction_models/PyTorchMLPRegressor.py
+++ b/freqtrade/freqai/prediction_models/PyTorchMLPRegressor.py
@@ -27,9 +27,9 @@ class PyTorchMLPRegressor(BasePyTorchRegressor):
             "model_training_parameters" : {
                 "learning_rate": 3e-4,
                 "trainer_kwargs": {
-                    "max_iters": 5000,
+                    "n_steps": 5000,
                     "batch_size": 64,
-                    "max_n_eval_batches": null,
+                    "n_epochs": null,
                 },
                 "model_kwargs": {
                     "hidden_dim": 512,
diff --git a/freqtrade/freqai/prediction_models/PyTorchTransformerRegressor.py b/freqtrade/freqai/prediction_models/PyTorchTransformerRegressor.py
index a76bab05c..846d6df2e 100644
--- a/freqtrade/freqai/prediction_models/PyTorchTransformerRegressor.py
+++ b/freqtrade/freqai/prediction_models/PyTorchTransformerRegressor.py
@@ -30,9 +30,9 @@ class PyTorchTransformerRegressor(BasePyTorchRegressor):
             "model_training_parameters" : {
                 "learning_rate": 3e-4,
                 "trainer_kwargs": {
-                    "max_iters": 5000,
+                    "n_steps": 5000,
                     "batch_size": 64,
-                    "max_n_eval_batches": null
+                    "n_epochs": null
                 },
                 "model_kwargs": {
                     "hidden_dim": 512,
diff --git a/freqtrade/freqai/torch/PyTorchDataConvertor.py b/freqtrade/freqai/torch/PyTorchDataConvertor.py
index e6b815373..0af14dd14 100644
--- a/freqtrade/freqai/torch/PyTorchDataConvertor.py
+++ b/freqtrade/freqai/torch/PyTorchDataConvertor.py
@@ -1,5 +1,4 @@
 from abc import ABC, abstractmethod
-from typing import Optional
 
 import pandas as pd
 import torch
@@ -12,14 +11,14 @@ class PyTorchDataConvertor(ABC):
     """
 
     @abstractmethod
-    def convert_x(self, df: pd.DataFrame, device: Optional[str] = None) -> torch.Tensor:
+    def convert_x(self, df: pd.DataFrame, device: str) -> torch.Tensor:
         """
         :param df: "*_features" dataframe.
         :param device: The device to use for training (e.g. 'cpu', 'cuda').
         """
 
     @abstractmethod
-    def convert_y(self, df: pd.DataFrame, device: Optional[str] = None) -> torch.Tensor:
+    def convert_y(self, df: pd.DataFrame, device: str) -> torch.Tensor:
         """
         :param df: "*_labels" dataframe.
         :param device: The device to use for training (e.g. 'cpu', 'cuda').
@@ -33,8 +32,8 @@ class DefaultPyTorchDataConvertor(PyTorchDataConvertor):
 
     def __init__(
             self,
-            target_tensor_type: Optional[torch.dtype] = None,
-            squeeze_target_tensor: bool = False
+            target_tensor_type: torch.dtype = torch.float32,
+            squeeze_target_tensor: bool = False,
     ):
         """
         :param target_tensor_type: type of target tensor, for classification use
@@ -45,23 +44,14 @@ class DefaultPyTorchDataConvertor(PyTorchDataConvertor):
         self._target_tensor_type = target_tensor_type
         self._squeeze_target_tensor = squeeze_target_tensor
 
-    def convert_x(self, df: pd.DataFrame, device: Optional[str] = None) -> torch.Tensor:
-        x = torch.from_numpy(df.values).float()
-        if device:
-            x = x.to(device)
-
+    def convert_x(self, df: pd.DataFrame, device: str) -> torch.Tensor:
+        numpy_arrays = df.values
+        x = torch.tensor(numpy_arrays, device=device, dtype=torch.float32)
         return x
 
-    def convert_y(self, df: pd.DataFrame, device: Optional[str] = None) -> torch.Tensor:
-        y = torch.from_numpy(df.values)
-
-        if self._target_tensor_type:
-            y = y.to(self._target_tensor_type)
-
+    def convert_y(self, df: pd.DataFrame, device: str) -> torch.Tensor:
+        numpy_arrays = df.values
+        y = torch.tensor(numpy_arrays, device=device, dtype=self._target_tensor_type)
         if self._squeeze_target_tensor:
             y = y.squeeze()
-
-        if device:
-            y = y.to(device)
-
         return y
diff --git a/freqtrade/freqai/torch/PyTorchModelTrainer.py b/freqtrade/freqai/torch/PyTorchModelTrainer.py
index 603e7ac12..371a953e7 100644
--- a/freqtrade/freqai/torch/PyTorchModelTrainer.py
+++ b/freqtrade/freqai/torch/PyTorchModelTrainer.py
@@ -1,5 +1,4 @@
 import logging
-import math
 from pathlib import Path
 from typing import Any, Dict, List, Optional
 
@@ -40,23 +39,27 @@ class PyTorchModelTrainer(PyTorchTrainerInterface):
             state_dict and model_meta_data saved by self.save() method.
         :param model_meta_data: Additional metadata about the model (optional).
         :param data_convertor: convertor from pd.DataFrame to torch.tensor.
-        :param max_iters: The number of training iterations to run.
-            iteration here refers to the number of times we call
-            self.optimizer.step(). used to calculate n_epochs.
+        :param n_steps: used to calculate n_epochs. The number of training iterations to run.
+            iteration here refers to the number of times optimizer.step() is called.
+            ignored if n_epochs is set.
+        :param n_epochs: The maximum number batches to use for evaluation.
         :param batch_size: The size of the batches to use during training.
-        :param max_n_eval_batches: The maximum number batches to use for evaluation.
         """
         self.model = model
         self.optimizer = optimizer
         self.criterion = criterion
         self.model_meta_data = model_meta_data
         self.device = device
-        self.max_iters: int = kwargs.get("max_iters", 100)
+        self.n_epochs: Optional[int] = kwargs.get("n_epochs", 10)
+        self.n_steps: Optional[int] = kwargs.get("n_steps", None)
+        if self.n_steps is None and not self.n_epochs:
+            raise Exception("Either `n_steps` or `n_epochs` should be set.")
+
         self.batch_size: int = kwargs.get("batch_size", 64)
-        self.max_n_eval_batches: Optional[int] = kwargs.get("max_n_eval_batches", None)
         self.data_convertor = data_convertor
         self.window_size: int = window_size
         self.tb_logger = tb_logger
+        self.test_batch_counter = 0
 
     def fit(self, data_dictionary: Dict[str, pd.DataFrame], splits: List[str]):
         """
@@ -72,55 +75,46 @@ class PyTorchModelTrainer(PyTorchTrainerInterface):
            backpropagation.
          - Updates the model's parameters using an optimizer.
         """
-        data_loaders_dictionary = self.create_data_loaders_dictionary(data_dictionary, splits)
-        epochs = self.calc_n_epochs(
-            n_obs=len(data_dictionary["train_features"]),
-            batch_size=self.batch_size,
-            n_iters=self.max_iters
-        )
         self.model.train()
-        for epoch in range(1, epochs + 1):
-            for i, batch_data in enumerate(data_loaders_dictionary["train"]):
 
+        data_loaders_dictionary = self.create_data_loaders_dictionary(data_dictionary, splits)
+        n_obs = len(data_dictionary["train_features"])
+        n_epochs = self.n_epochs or self.calc_n_epochs(n_obs=n_obs)
+        batch_counter = 0
+        for _ in range(n_epochs):
+            for _, batch_data in enumerate(data_loaders_dictionary["train"]):
                 xb, yb = batch_data
-                xb.to(self.device)
-                yb.to(self.device)
+                xb = xb.to(self.device)
+                yb = yb.to(self.device)
                 yb_pred = self.model(xb)
                 loss = self.criterion(yb_pred, yb)
 
                 self.optimizer.zero_grad(set_to_none=True)
                 loss.backward()
                 self.optimizer.step()
-                self.tb_logger.log_scalar("train_loss", loss.item(), i)
+                self.tb_logger.log_scalar("train_loss", loss.item(), batch_counter)
+                batch_counter += 1
 
             # evaluation
             if "test" in splits:
-                self.estimate_loss(
-                    data_loaders_dictionary,
-                    self.max_n_eval_batches,
-                    "test"
-                )
+                self.estimate_loss(data_loaders_dictionary, "test")
 
     @torch.no_grad()
     def estimate_loss(
             self,
             data_loader_dictionary: Dict[str, DataLoader],
-            max_n_eval_batches: Optional[int],
             split: str,
     ) -> None:
         self.model.eval()
-        n_batches = 0
-        for i, batch_data in enumerate(data_loader_dictionary[split]):
-            if max_n_eval_batches and i > max_n_eval_batches:
-                n_batches += 1
-                break
+        for _, batch_data in enumerate(data_loader_dictionary[split]):
             xb, yb = batch_data
-            xb.to(self.device)
-            yb.to(self.device)
+            xb = xb.to(self.device)
+            yb = yb.to(self.device)
 
             yb_pred = self.model(xb)
             loss = self.criterion(yb_pred, yb)
-            self.tb_logger.log_scalar(f"{split}_loss", loss.item(), i)
+            self.tb_logger.log_scalar(f"{split}_loss", loss.item(), self.test_batch_counter)
+            self.test_batch_counter += 1
 
         self.model.train()
 
@@ -148,31 +142,30 @@ class PyTorchModelTrainer(PyTorchTrainerInterface):
 
         return data_loader_dictionary
 
-    @staticmethod
-    def calc_n_epochs(n_obs: int, batch_size: int, n_iters: int) -> int:
+    def calc_n_epochs(self, n_obs: int) -> int:
         """
         Calculates the number of epochs required to reach the maximum number
         of iterations specified in the model training parameters.
 
-        the motivation here is that `max_iters` is easier to optimize and keep stable,
+        the motivation here is that `n_steps` is easier to optimize and keep stable,
         across different n_obs - the number of data points.
         """
+        assert isinstance(self.n_steps, int), "Either `n_steps` or `n_epochs` should be set."
+        n_batches = n_obs // self.batch_size
+        n_epochs = min(self.n_steps // n_batches, 1)
+        if n_epochs <= 10:
+            logger.warning(
+                f"Setting low n_epochs: {n_epochs}. "
+                f"Please consider increasing `n_steps` hyper-parameter."
+            )
 
-        n_batches = math.ceil(n_obs // batch_size)
-        epochs = math.ceil(n_iters // n_batches)
-        if epochs <= 10:
-            logger.warning("User set `max_iters` in such a way that the trainer will only perform "
-                           f" {epochs} epochs. Please consider increasing this value accordingly")
-            if epochs <= 1:
-                logger.warning("Epochs set to 1. Please review your `max_iters` value")
-                epochs = 1
-        return epochs
+        return n_epochs
 
     def save(self, path: Path):
         """
         - Saving any nn.Module state_dict
         - Saving model_meta_data, this dict should contain any additional data that the
-          user needs to store. e.g class_names for classification models.
+          user needs to store. e.g. class_names for classification models.
         """
 
         torch.save({
diff --git a/tests/freqai/conftest.py b/tests/freqai/conftest.py
index 4c4891ceb..9c7a950e7 100644
--- a/tests/freqai/conftest.py
+++ b/tests/freqai/conftest.py
@@ -97,9 +97,9 @@ def mock_pytorch_mlp_model_training_parameters() -> Dict[str, Any]:
     return {
             "learning_rate": 3e-4,
             "trainer_kwargs": {
-                "max_iters": 1,
+                "n_steps": None,
                 "batch_size": 64,
-                "max_n_eval_batches": 1,
+                "n_epochs": 1,
             },
             "model_kwargs": {
                 "hidden_dim": 32,