From ad695a818b0c213f57e2e3dd28794b5915026ce9 Mon Sep 17 00:00:00 2001
From: Sven Mika <sven@anyscale.io>
Date: Sat, 6 Jun 2020 22:47:42 +0200
Subject: [PATCH] Bug fix in the contextual bandit's linear_regression.py
 model. (#8815)

---
 rllib/contrib/bandits/models/linear_regression.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/rllib/contrib/bandits/models/linear_regression.py b/rllib/contrib/bandits/models/linear_regression.py
index 568fcd56d..76ff2203a 100644
--- a/rllib/contrib/bandits/models/linear_regression.py
+++ b/rllib/contrib/bandits/models/linear_regression.py
@@ -34,7 +34,7 @@ class OnlineLinearRegression(nn.Module):
     def partial_fit(self, x, y):
         # TODO: Handle batch of data rather than individual points
         self._check_inputs(x, y)
-        x = x.squeeze()
+        x = x.squeeze(0)
         y = y.item()
         self.time += 1
         self.delta_f += y * x