support dict batch in dynabert (#1064)

paddleslim/nas/ofa/utils/nlp_utils.py

@@ -70,6 +70,10 @@ def compute_neuron_head_importance(task_name,
         data_loader = (data_loader, )
     for data in data_loader:
         for batch in data:
+            if isinstance(batch, dict):
+                input_ids, segment_ids, labels = batch['input_ids'], batch[
+                    'token_type_ids'], batch['labels']
+            else:
                 input_ids, segment_ids, labels = batch
             logits = model(
                 input_ids, segment_ids, attention_mask=[None, head_mask])