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