Fix input dim (#365)
This commit is contained in:
@@ -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.
|
# 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(
|
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_feature_projector = MLP(
|
||||||
self.rgb_encoder.feature_dim, hidden_channels=[self.config.gpt_input_dim]
|
self.rgb_encoder.feature_dim, hidden_channels=[self.config.gpt_input_dim]
|
||||||
|
|||||||
Reference in New Issue
Block a user