fixing tests

    restore: bool

      if True, restore the model from the most recent checkpoint and continue training

      from there.  If False, retrain the model from scratch.

    variables: list of torch.nn.Parameter

    variables: list of hk.Variable

      the variables to train.  If None (the default), all trainable variables in

      the model are used.

    def __init__(self, output_size: int = 1):

      super().__init__()

      self._network = hk.Sequential([hk.Linear(200), jax.nn.relu])

      self._network1 = hk.Sequential([hk.Linear(200), jax.nn.relu])

      self._network2 = hk.Sequential([hk.Linear(200), jax.nn.relu])

      self.output = hk.Linear(output_size)

      self.log_var = hk.Linear(output_size)

    def __call__(self, x):

      # x, dropout_rate = x

      x = self._network(x)

      # if x is not None:

      x = self._network1(x)

      x = hk.dropout(hk.next_rng_key(), 0.1, x)

      x = self._network2(x)

      x = hk.dropout(hk.next_rng_key(), 0.1, x)

      output = self.output(x)

      log_var = self.log_var(x)

      learning_rate=0.003)

  model.fit(dataset, nb_epochs=2500)

  pred, std = model.predict_uncertainty(dataset)

  assert np.mean(np.abs(y - pred)) < 1.0

  assert np.mean(np.abs(y - pred)) < 2.0

  assert noise < np.mean(std) < 1.0