code: def infer(self, x, **fit_params):
          
          x = to_tensor(x, device=self.device)
          if isinstance(x, dict):
              x_dict = self._merge_x_and_fit_params(x, fit_params)
              return self.module_(**x_dict)
          return self.module_(x, **fit_params)
docs: Perform a single inference step on a batch of data.
  
          Parameters
          ----------
          x : input data
            A batch of the input data.
  
          **fit_params : dict
            Additional parameters passed to the ``forward`` method of
            the module and to the ``self.train_split`` call.