snapshot...

modules/deeplearning/cloud_opd_srcnn.py

@@ -14,7 +14,7 @@ LOG_DEVICE_PLACEMENT = False
 PROC_BATCH_SIZE = 4
 PROC_BATCH_BUFFER_SIZE = 5000
 
-NumClasses = 5
+NumClasses = 2
 if NumClasses == 2:
     NumLogits = 1
 else:
@@ -499,7 +499,6 @@ class SRCNN:
 
     @tf.function(input_signature=[tf.TensorSpec(None, tf.float32), tf.TensorSpec(None, tf.float32)])
     def train_step(self, inputs, labels):
-        labels = tf.squeeze(labels, axis=[3])
         with tf.GradientTape() as tape:
             pred = self.model([inputs], training=True)
             loss = self.loss(labels, pred)
@@ -519,7 +518,6 @@ class SRCNN:
 
     @tf.function(input_signature=[tf.TensorSpec(None, tf.float32), tf.TensorSpec(None, tf.float32)])
     def test_step(self, inputs, labels):
-        labels = tf.squeeze(labels, axis=[3])
         pred = self.model([inputs], training=False)
         t_loss = self.loss(labels, pred)