def r_squared(yhatss, yss, return_error=False):
    """
  Returns the coefficient of determination of two 2d tensors
  (where the first dimension in each indexes the examples),
  one holding the yhatss (the predicted outputs) and the other
  holding the actual outputs, yss.

  Note: this is rigorously relevant only to linear, meaning
  possibly polynomial linear, regression.

  Args:
    yhatss (torch.Tensor): The predicted outputs.
    yss (torch.Tensor): The actual outputs.
    return_error (bool): If False return the proportion of the
        variation explained by the regression line. If True,
        return 1 minus that proportion. Default: False.
  """
    SS_E, SS_T = 0.0, 0.0
    mean = yss.mean(0)
    for yhats, ys in zip(yhatss, yss):
        SS_E += (ys - yhats).pow(2)
        SS_T += (ys - mean).pow(2)
    if return_error:
        return (SS_E / SS_T).item()
    else:
        return 1.0 - (SS_E / SS_T).item()
