reverse reshaping handler output

handler.py

@@ -47,7 +47,5 @@ class EndpointHandler:
         output = self.model(**tokenized_inputs)
         predictions = torch.argmax(output.logits, dim=-1).numpy(force=True)
         batch['label'] = [self.model.config.id2label[p] for p in predictions]
-        
-        n = len(test_dict['id'])
-        result = [{key : value[i] for key, value in test_dict.items() if key not in ('text')} for i in range(n)]
-        return result
+        batch.pop('text')
+        return batch