diff --git a/modules/deeplearning/espcn.py b/modules/deeplearning/espcn.py
index 4086e00626256bd63cfb3cb62e6ba8a658618b3c..2f02c9a268b94ca3e3e91891b93466131b959e51 100644
--- a/modules/deeplearning/espcn.py
+++ b/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)