    def predict(self, inputs, factored=False, *args, **kwargs):
        """Returns the distribution predicted by the model for each input vector in inputs.
        Behavior is affected by the dimensionality of inputs and factored as follows:

        inputs is 2D, factored=True: Each row is treated as an input vector.
            Returns a mean of shape [ensemble_size, batch_size, output_dim] and variance of shape
            [ensemble_size, batch_size, output_dim], where N(mean[i, j, :], diag([i, j, :])) is the
            predicted output distribution by the ith model in the ensemble on input vector j.

        inputs is 2D, factored=False: Each row is treated as an input vector.
            Returns a mean of shape [batch_size, output_dim] and variance of shape
            [batch_size, output_dim], where aggregation is performed as described in the paper.

        inputs is 3D, factored=True/False: Each row in the last dimension is treated as an input vector.
            Returns a mean of shape [ensemble_size, batch_size, output_dim] and variance of sha
            [ensemble_size, batch_size, output_dim], where N(mean[i, j, :], diag([i, j, :])) is the
            predicted output distribution by the ith model in the ensemble on input vector [i, j].

        Arguments:
            inputs (np.ndarray): An array of input vectors in rows. See above for behavior.
            factored (bool): See above for behavior.
        """
        if len(inputs.shape) == 2:
            if factored:
                return self.sess.run(
                    [self.sy_pred_mean2d_fac, self.sy_pred_var2d_fac],
                    feed_dict={self.sy_pred_in2d: inputs}
                )
            else:
                return self.sess.run(
                    [self.sy_pred_mean2d, self.sy_pred_var2d],
                    feed_dict={self.sy_pred_in2d: inputs}
                )
        else:
            return self.sess.run(
                [self.sy_pred_mean3d_fac, self.sy_pred_var3d_fac],
                feed_dict={self.sy_pred_in3d: inputs}
            )
Returns the distribution predicted by the model for each input vector in inputs.
Behavior is affected by the dimensionality of inputs and factored as follows:

inputs is 2D, factored=True: Each row is treated as an input vector.
Returns a mean of shape [ensemble_size, batch_size, output_dim] and variance of shape
[ensemble_size, batch_size, output_dim], where N(mean[i, j, :], diag([i, j, :])) is the
predicted output distribution by the ith model in the ensemble on input vector j.

inputs is 2D, factored=False: Each row is treated as an input vector.
Returns a mean of shape [batch_size, output_dim] and variance of shape
[batch_size, output_dim], where aggregation is performed as described in the paper.

inputs is 3D, factored=True/False: Each row in the last dimension is treated as an input vector.
Returns a mean of shape [ensemble_size, batch_size, output_dim] and variance of sha
[ensemble_size, batch_size, output_dim], where N(mean[i, j, :], diag([i, j, :])) is the
predicted output distribution by the ith model in the ensemble on input vector [i, j].

Arguments:
inputs (np.ndarray): An array of input vectors in rows. See above for behavior.
factored (bool): See above for behavior.
mbpo/models/bnn.py