From 0f6e0f6d74f1c26102c13113d3c21227078adedd Mon Sep 17 00:00:00 2001 From: Alexander Soare Date: Mon, 19 Aug 2024 11:42:32 +0100 Subject: [PATCH] Fix input dim (#365) --- lerobot/common/policies/vqbet/modeling_vqbet.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lerobot/common/policies/vqbet/modeling_vqbet.py b/lerobot/common/policies/vqbet/modeling_vqbet.py index bc12dfa2..a73acb4f 100644 --- a/lerobot/common/policies/vqbet/modeling_vqbet.py +++ b/lerobot/common/policies/vqbet/modeling_vqbet.py @@ -289,7 +289,7 @@ class VQBeTModel(nn.Module): # To input state and observation features into GPT layers, we first project the features to fit the shape of input size of GPT. self.state_projector = MLP( - config.output_shapes["action"][0], hidden_channels=[self.config.gpt_input_dim] + config.input_shapes["observation.state"][0], hidden_channels=[self.config.gpt_input_dim] ) self.rgb_feature_projector = MLP( self.rgb_encoder.feature_dim, hidden_channels=[self.config.gpt_input_dim]