diff --git a/n3fit/src/n3fit/layers/x_operations.py b/n3fit/src/n3fit/layers/x_operations.py
index a55b56d2ba..3ed9705b59 100644
--- a/n3fit/src/n3fit/layers/x_operations.py
+++ b/n3fit/src/n3fit/layers/x_operations.py
@@ -76,3 +76,6 @@ def __init__(self, grid_weights, x_axis=2, **kwargs):
 
     def call(self, pdf):
         return op.tensor_product(pdf, self.grid_weights, axes=[self.x_axis, 0])
+
+    def compute_output_shape(self, input_shape):
+        return input_shape[: self.x_axis] + input_shape[self.x_axis + 1 :]
diff --git a/n3fit/src/n3fit/msr.py b/n3fit/src/n3fit/msr.py
index e781a672d2..1bca92f2e1 100644
--- a/n3fit/src/n3fit/msr.py
+++ b/n3fit/src/n3fit/msr.py
@@ -85,7 +85,7 @@ def generate_msr_model_and_grid(
     )([x_divided, pdf_xgrid_integration])
 
     # 4. Integrate the pdf
-    pdf_integrated = xIntegrator(weights_array, input_shape=(nx,))(pdf_integrand)
+    pdf_integrated = xIntegrator(weights_array, input_shape=(1, replicas, nx))(pdf_integrand)
 
     # 5. THe input for the photon integral, will be set to 0 if no photons
     photon_integral = Input(shape=(replicas, 1), batch_size=1, name='photon_integral')