[MEGA] nit size test (huggingface#23028)

ArthurZucker · web-flow · commit 1933231a0a98 · 2023-04-27T16:21:00.000+02:00
* add fast not use warning

* properly check sequence_length vs chunk_size

* fixup
diff --git a/src/transformers/models/mega/modeling_mega.py b/src/transformers/models/mega/modeling_mega.py
@@ -1541,12 +1541,6 @@ def forward(
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        if self.config.use_chunking and (input_ids.size(1) > self.config.chunk_size):
-            if input_ids.size(1) % self.config.chunk_size != 0:
-                raise ValueError(
-                    f"config.use_chunking is activated; input sequence length must be shorter than or a multiple of config.chunk_size\nreceived sequence length of {input_ids.size(1)} with chunk size {self.config.chunk_size}"
-                )
-
         if input_ids is not None and inputs_embeds is not None:
             raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
         elif input_ids is not None:
@@ -1560,6 +1554,12 @@ def forward(
 
         batch_size, sequence_length = input_shape
 
+        if self.config.use_chunking and (sequence_length > self.config.chunk_size):
+            if sequence_length % self.config.chunk_size != 0:
+                raise ValueError(
+                    f"config.use_chunking is activated; input sequence length must be shorter than or a multiple of config.chunk_size\nreceived sequence length of {sequence_length} with chunk size {self.config.chunk_size}"
+                )
+
         if self.config.is_decoder:
             use_cache = use_cache if use_cache is not None else self.config.use_cache
 
