workaround for squeeze bug in pytorch
@@ -65,7 +65,7 @@ def model(data):
x_data = data[:, :-1]
y_data = data[:, -1]
# run the regressor forward conditioned on inputs
- prediction_mean = lifted_reg_model(x_data).squeeze()
+ prediction_mean = lifted_reg_model(x_data).squeeze(1)
pyro.sample("obs",
Normal(prediction_mean, Variable(torch.ones(data.size(0))).type_as(data)),
obs=y_data.squeeze())
