Removed unnecessary `self.projection` call in `VivitTubeletEmbeddings` (#31632)

removes unnecessary second projection call

Removed unnecessary `self.projection` call in `VivitTubeletEmbeddings` (#31632)
removes unnecessary second projection call
1218e439 · Vladimir Iashin · GitHub · 2daf2c3e · 1218e439
Unverified Commit 1218e439 authored Jun 26, 2024 by Vladimir Iashin Committed by GitHub Jun 26, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

src/transformers/models/vivit/modeling_vivit.py src/transformers/models/vivit/modeling_vivit.py +2 -1

No files found.
--- a/src/transformers/models/vivit/modeling_vivit.py
+++ b/src/transformers/models/vivit/modeling_vivit.py
@@ -76,7 +76,8 @@ class VivitTubeletEmbeddings(nn.Module):

        x = self.projection(pixel_values)
        # out_batch_size, out_num_channels, out_num_frames, out_height, out_width = x.shape
-        x = self.projection(pixel_values).flatten(2).transpose(1, 2)
+        # flattens time and space dimensions, transposes to (out_batch_size, flat_tokens, out_num_channels)
+        x = x.flatten(2).transpose(1, 2)
        return x