<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def train_step(self, Xi, yi, **fit_params):
 """Prepares a loss function callable and pass it to the optimizer, hence performing one optimization step. Loss function callable as required by some optimizers (and accepted by all of them):
 https://pytorch.org/docs/master/optim.html#optimizer-step-closure The module is set to be in train mode (e.g. dropout is applied). Parameters Xi : input data A batch of the input data. yi : target data A batch of the target data. **fit_params : dict Additional parameters passed to the ``forward`` method of the module and to the train_split call. """

        step_accumulator = self.get_train_step_accumulator()
        def step_fn():
            step = self.train_step_single(Xi, yi, **fit_params)
            step_accumulator.store_step(step)
            return step['loss']
        self.optimizer_.step(step_fn)
        return step_accumulator.get_step()