condition strategy

broadinstitute · Dec 16, 2024 · d85a97d · d85a97d
1 parent 82a9d67
commit d85a97d
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/ml4h/models/diffusion_blocks.py b/ml4h/models/diffusion_blocks.py
@@ -108,9 +108,12 @@ 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, 1, 1, filters))
-    beta = tf.reshape(beta, (-1, 1, 1, filters))
-
+    if 4 == len(input_tensor.shape):
+        gamma = tf.reshape(gamma, (-1, 1, 1, filters))
+        beta = tf.reshape(beta, (-1, 1, 1, filters))
+    elif 3 == len(input_tensor.shape):
+        gamma = tf.reshape(gamma, (-1, 1, filters))
+        beta = tf.reshape(beta, (-1, 1, filters))
     # Apply FiLM (Feature-wise Linear Modulation)
     return input_tensor * gamma + beta