snapshot...

modules/deeplearning/espcn.py

@@ -451,7 +451,7 @@ class ESPCN:
             activation = tf.nn.softmax  # For multi-class
 
         # Called logits, but these are actually probabilities, see activation
-        self.logits = tf.keras.layers.Conv2D(1, kernel_size=1, strides=1, padding=padding, name='probability', activation=activation)(conv)
+        self.logits = tf.keras.layers.Activation(activation=activation)(conv)
 
         print(self.logits.shape)