Add output_dim parameter to LinearRegression (#198)

pl_bolts/models/regression/linear_regression.py

@@ -14,6 +14,7 @@ class LinearRegression(pl.LightningModule):
 
  def __init__(self,
  input_dim: int,
+ output_dim: int = 1,
  bias: bool = True,
  learning_rate: float = 0.0001,
  optimizer: Optimizer = Adam,
@@ -26,6 +27,7 @@ def __init__(self,
 
  Args:
  input_dim: number of dimensions of the input (1+)
+ output_dim: number of dimensions of the output (default=1)
  bias: If false, will not use $+b$
  learning_rate: learning_rate for the optimizer
  optimizer: the optimizer to use (default='Adam')
@@ -37,7 +39,7 @@ def __init__(self,
  self.save_hyperparameters()
  self.optimizer = optimizer
 
- self.linear = nn.Linear(in_features=self.hparams.input_dim, out_features=1, bias=bias)
+ self.linear = nn.Linear(in_features=self.hparams.input_dim, out_features=self.hparams.output_dim, bias=bias)
 
  def forward(self, x):
  y_hat = self.linear(x)
@@ -114,6 +116,7 @@ def add_model_specific_args(parent_parser):
  parser = ArgumentParser(parents=[parent_parser], add_help=False)
  parser.add_argument('--learning_rate', type=float, default=0.0001)
  parser.add_argument('--input_dim', type=int, default=None)
+ parser.add_argument('--output_dim', type=int, default=1)
  parser.add_argument('--bias', default='store_true')
  parser.add_argument('--batch_size', type=int, default=16)
  return parser