diff --git a/ml4h/models/diffusion_blocks.py b/ml4h/models/diffusion_blocks.py
index 6cd5835be..2ac518d62 100644
--- a/ml4h/models/diffusion_blocks.py
+++ b/ml4h/models/diffusion_blocks.py
@@ -106,8 +106,8 @@ def condition_layer_film(input_tensor, control_vector, filters):
     beta = layers.Dense(filters, activation="linear")(control_vector)
 
     # Reshape gamma and beta to match the spatial dimensions
-    gamma = tf.reshape(gamma, (-1,) + input_tensor.shape[1:-1] + (filters,))
-    beta = tf.reshape(beta, (-1,) + input_tensor.shape[1:-1] + (filters,))
+    #gamma = tf.reshape(gamma, (-1,) + input_tensor.shape[1:-1] + (filters,))
+    #beta = tf.reshape(beta, (-1,) + input_tensor.shape[1:-1] + (filters,))
     # Apply FiLM (Feature-wise Linear Modulation)
     return input_tensor * gamma + beta