[Reformer] Axial Pos Emb Improve mem usage reformer (#5209)

* improve mem handling

* improve mem for pos ax encodings

src/transformers/modeling_reformer.py

@@ -154,9 +154,14 @@ class AxialPositionEmbeddings(nn.Module):
                 self.axial_pos_shape, sequence_length, self.least_common_mult_chunk_length,
             )
 
-            # reshape axial encodings and use only until sequence_length
-            position_encodings = torch.cat(broadcasted_weights, dim=-1)
-            position_encodings = position_encodings.view(batch_size, -1, position_encodings.shape[-1])[
+            # compute how many columns are needed
+            required_pos_encodings_columns = -(-sequence_length // self.axial_pos_shape[1])
+
+            # cut to columns that are needed
+            position_encodings = torch.cat(
+                [weight[:, :required_pos_encodings_columns] for weight in broadcasted_weights], dim=-1
+            )
+            position_encodings = torch.reshape(position_encodings, (batch_size, -1, position_encodings.shape[-1]))[
                 :, :sequence_length
             ]