minor: nit

This commit is contained in:
jayLEE0301 2024-06-04 20:45:26 -04:00
parent 6d72847bfe
commit c4c5977f37
1 changed files with 1 additions and 3 deletions

View File

@ -314,9 +314,7 @@ class VQBeTModel(nn.Module):
features[:, -len_additional_action_token:]
], dim=1)
# pass through action head
pred_action = self.action_head(
features,
)
pred_action = self.action_head(features)
# if rollout, VQ-BeT don't calculate loss
if rollout:
return pred_action["predicted_action"][:, n_obs_steps-1, :].reshape(batch_size, self.config.action_chunk_size, -1)