From c4c222a3813791f07776e6357512f953006d0138 Mon Sep 17 00:00:00 2001 From: Dibya Ghosh Date: Mon, 15 Apr 2024 11:59:44 -0700 Subject: [PATCH] Fix potential off-by-one error in attention mask generation --- octo/model/components/block_transformer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/octo/model/components/block_transformer.py b/octo/model/components/block_transformer.py index c4107e35..851195b6 100644 --- a/octo/model/components/block_transformer.py +++ b/octo/model/components/block_transformer.py @@ -290,7 +290,7 @@ def generate_attention_mask( self.verify_causality(prefix_groups, timestep_groups) def _get_position(i, tokens_per_elem): - return np.searchsorted(np.cumsum(tokens_per_elem), i) + return np.searchsorted(np.cumsum(tokens_per_elem), i, side='right') horizon = timestep_groups[0].tokens.shape[1] tokens_per_prefix_group = [group.tokens.shape[1] for group in prefix_groups]