From dd4155359437876464342e71b918659c2222ec55 Mon Sep 17 00:00:00 2001 From: Cao Hanzhe Date: Fri, 8 May 2026 23:51:17 +0800 Subject: [PATCH] fix: squeeze 3D action tensor in LinUCB learn_batch batch.action can have shape [B, 1, N] for one-hot encoded actions, but torch.cat with batch.state (shape [B, D]) requires 2D tensors. Squeeze dim=1 to handle both [B, N] and [B, 1, N] action shapes. Fixes #125 --- pearl/policy_learners/contextual_bandits/linear_bandit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pearl/policy_learners/contextual_bandits/linear_bandit.py b/pearl/policy_learners/contextual_bandits/linear_bandit.py index ff0e1cd8..606011d8 100644 --- a/pearl/policy_learners/contextual_bandits/linear_bandit.py +++ b/pearl/policy_learners/contextual_bandits/linear_bandit.py @@ -150,7 +150,7 @@ def learn_batch(self, batch: TransitionBatch) -> dict[str, Any]: if batch.weight is not None else torch.ones_like(expected_values) ) - x = torch.cat([batch.state, batch.action], dim=1) + x = torch.cat([batch.state, torch.squeeze(batch.action, dim=1)], dim=1) self.model.learn_batch( x=x, y=batch.reward,