From ddfdf9aa760cade1144948ea4832533d4b572aa2 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Thu, 8 Jan 2026 14:44:09 +0100 Subject: [PATCH] unsqueeze after unnromalize --- examples/openarms/evaluate_relative.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/examples/openarms/evaluate_relative.py b/examples/openarms/evaluate_relative.py index 3c708c0e8..b403faac5 100644 --- a/examples/openarms/evaluate_relative.py +++ b/examples/openarms/evaluate_relative.py @@ -156,7 +156,11 @@ def inference_loop_relative( elif action_tensor.dim() == 2: action_tensor = action_tensor.unsqueeze(1) # [batch, 1, action_dim] action_tensor = relative_normalizer.unnormalize(action_tensor) - action_tensor = action_tensor.squeeze(1) # back to [batch, action_dim] + action_tensor = action_tensor.squeeze() # remove all size-1 dims -> [action_dim] + + # Ensure tensor is 1D or 2D with batch=1 for conversion + while action_tensor.dim() > 2: + action_tensor = action_tensor.squeeze(0) # Convert tensor to dict relative_action = tensor_to_robot_action(action_tensor, dataset.features)