def test_spawned_streamer(self):
# Spawn releases 4 gpu worker processes
streamer = Streamer(self.vision_model.batch_prediction, batch_size=8, worker_num=4, cuda_devices=(0, 1, 2, 3))
single_predict = streamer.predict(self.input_batch)
assert single_predict == self.single_output
batch_predict = streamer.predict(self.input_batch * BATCH_SIZE)
assert batch_predict == self.batch_output
def test_spawned_streamer():
# Spawn releases 4 gpu worker processes
streamer = Streamer(vision_model.batch_prediction, batch_size=8, worker_num=4)
single_predict = streamer.predict(input_batch)
assert single_predict == single_output
batch_predict = streamer.predict(input_batch * BATCH_SIZE)
assert batch_predict == batch_output