From 9313c8bb7f62b14f8c15d2119ae678641f751dbd Mon Sep 17 00:00:00 2001 From: Alex Auvolat Date: Thu, 23 Jul 2015 10:13:28 -0400 Subject: Can do test with batches of size >1 --- ext_test.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) (limited to 'ext_test.py') diff --git a/ext_test.py b/ext_test.py index 3af637b..9e64223 100644 --- a/ext_test.py +++ b/ext_test.py @@ -61,12 +61,13 @@ class RunOnTest(SimpleExtension): for d in self.test_stream.get_epoch_iterator(as_dict=True): input_values = [d[k.name] for k in self.inputs] output_values = self.function(*input_values) - if 'destination' in self.outputs: - destination = output_values[self.outputs.index('destination')] - dest_outcsv.writerow([d['trip_id'][0], destination[0, 0], destination[0, 1]]) - if 'duration' in self.outputs: - duration = output_values[self.outputs.index('duration')] - time_outcsv.writerow([d['trip_id'][0], int(round(duration[0]))]) + for i in range(d['trip_id'].shape[0]): + if 'destination' in self.outputs: + destination = output_values[self.outputs.index('destination')] + dest_outcsv.writerow([d['trip_id'][i], destination[i, 0], destination[i, 1]]) + if 'duration' in self.outputs: + duration = output_values[self.outputs.index('duration')] + time_outcsv.writerow([d['trip_id'][i], int(round(duration[i]))]) if 'destination' in self.outputs: dest_outfile.close() -- cgit v1.2.3