hangX/lerobot
1
0
Fork 0
forked from tangger/lerobot
Code Pull Requests Activity

fix vqbet

Browse Source
This commit is contained in:
Remi Cadene
2024-07-17 00:45:34 +02:00
parent 4ea6481390
commit df303d7311
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
lerobot/common/policies/vqbet/modeling_vqbet.py
View File

@@ -298,7 +298,7 @@ class VQBeTModel(nn.Module):
# bin prediction head / offset prediction head part of VQ-BeT # bin prediction head / offset prediction head part of VQ-BeT
self.action_head = VQBeTHead(config) self.action_head = VQBeTHead(config)
num_tokens = self.config.n_action_pred_token + self.config.action_chunk_size - 1 num_tokens = self.config.n_action_pred_token
self.register_buffer( self.register_buffer(
"select_target_actions_indices", "select_target_actions_indices",
torch.row_stack([torch.arange(i, i + self.config.action_chunk_size) for i in range(num_tokens)]), torch.row_stack([torch.arange(i, i + self.config.action_chunk_size) for i in range(num_tokens)]),