fix attribution target (use mean)

dn070017 · Dec 14, 2023 · a2c9519 · a2c9519
1 parent 01e9c7c
commit a2c9519
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/cavachon/modules/parameterizers/IndependentZeroInflatedNegativeBinomial.py b/cavachon/modules/parameterizers/IndependentZeroInflatedNegativeBinomial.py
@@ -1,9 +1,11 @@
-from cavachon.environment.Constants import Constants
-from cavachon.modules.parameterizers.Parameterizer import Parameterizer
 from typing import Mapping
 
 import tensorflow as tf
 
+from cavachon.environment.Constants import Constants
+from cavachon.modules.parameterizers.Parameterizer import Parameterizer
+
+
 class IndependentZeroInflatedNegativeBinomial(Parameterizer):
   """IndependentZeroInflatedNegativeBinomial
 
@@ -66,7 +68,7 @@ def compute_attribution_target(self, inputs: tf.Tensor):
     outputs = self.layer(inputs.get(Constants.TENSOR_NAME_X))
     probs, means, dispersion = tf.split(outputs, 3, axis=-1)
     probs = tf.keras.activations.sigmoid(probs)
-    return tf.concat([probs, means, dispersion], axis=-1)
+    return means
 
   @classmethod
   def modify_outputs(