Fix bug in hierarchical pos embed (#51)

* fix: scholar link * fix: maxT in hierarchical_pos_embed
amazon-science · Jun 11, 2023 · 0303b37 · 0303b37
1 parent a9b42c1
commit 0303b37
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/earthformer/cuboid_transformer/cuboid_transformer.py b/src/earthformer/cuboid_transformer/cuboid_transformer.py
@@ -2365,7 +2365,7 @@ def __init__(self,
             if self.hierarchical_pos_embed:
                 self.hierarchical_pos_embed_l = nn.ModuleList([
                     PosEmbed(embed_dim=self.mem_shapes[i][-1], typ=pos_embed_type,
-                             maxT=self.mem_shapes[i][0], maxH=self.mem_shapes[i][1], maxW=self.mem_shapes[i][2])
+                             maxT=target_temporal_length, maxH=self.mem_shapes[i][1], maxW=self.mem_shapes[i][2])
                     for i in range(self.num_blocks - 1)])
 
         self.reset_parameters()