get example_input for QA from train dataset (intel#1091)

chensuyue · Jul 22, 2022 · 0b5a8da · 0b5a8da
1 parent ed5a997
commit 0b5a8da
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/neural_compressor/adaptor/pytorch.py b/neural_compressor/adaptor/pytorch.py
@@ -2508,6 +2508,8 @@ def get_example_inputs(self, dataloader):
             if "label" in example_inputs.keys():
                 example_inputs.pop("label")
             for key, value in example_inputs.items():
+                if key == "start_positions" or key == "end_positions":
+                    continue
                 input_tensor.append(value)
             return input_tensor
         if isinstance(example_inputs, list) or isinstance(example_inputs, tuple):