Update TEQ train dataloader (#1554)

changwangss · web-flow · commit d1e994becee1 · 2024-01-19T17:59:47.000+08:00
diff --git a/neural_compressor/adaptor/torch_utils/teq.py b/neural_compressor/adaptor/torch_utils/teq.py
@@ -256,7 +256,9 @@ def train(
 
         while global_steps <= train_steps:
             for inputs in dataloader:
-                if isinstance(inputs, dict):
+                if isinstance(inputs, torch.Tensor):
+                    input_id = inputs
+                elif isinstance(inputs, dict):
                     input_id = inputs["input_ids"]
                 else:
                     input_id = inputs[0]
