(128, 2, 224, 224, 3)
image a shape
(128, 224, 224, 3)
labels
tf.Tensor(0, shape=(), dtype=int32)
prediction a
tf.Tensor([ 0.00588202 -0.28380316  0.00119852 -0.33077663 -0.4034524 ], shape=(5,), dtype=float32)