def on_batch_end(self, runner: "IRunner") -> None:
        """On batch end action: get data from runner's batch
        and update a loader storage with it

        Args:
            runner: runner for the experiment.
        """
        assert (
            torch.isnan(runner.batch[self.feature_key]).sum() == 0
        ), "SklearnModelCallback can't process Tensors with NaN!"
        if runner.loader_key == self._train_loader:
            self.storage.update(**runner.batch)
        if runner.loader_key in self._valid_loaders:
            features = runner.batch[self.feature_key].detach().cpu().numpy()
            # classifier predict
            classifier_predict = getattr(self.model, self.predict_method)
            predictions = classifier_predict(features)
            runner.batch[self.predict_key] = torch.tensor(
                predictions, device=runner.engine.device
            )
On batch end action: get data from runner's batch
        and update a loader storage with it

        Args:
            runner: runner for the experiment.
        
catalyst/callbacks/sklearn_model.py