def test_model_broadcast(self): init_executor_gateway(self.sc) model = Linear(3, 2) broadcasted = broadcast_model(self.sc, model) input_data = np.random.rand(3) output = self.sc.parallelize([input_data], 1)\ .map(lambda x: broadcasted.value.forward(x)).first() expected = model.forward(input_data) assert_allclose(output, expected)