Why do we need to mask the digit capsule? We learned that in order to make sure that the network has learned the important features, we use a three-layer network called a decoder network, which tries to reconstruct the original image from the digit capsules. If the decoder is able to reconstruct the image successfully from the digit capsules, then it means the network has learned the important features of the image; otherwise, the network has not learned the correct features of the image.

The digit capsules contain the activity vector for all the digits. But the decoder wants to reconstruct only the given input digit (the input image). So, we mask out the activity vector of all the digits, except for the correct digit. Then we use this masked digit capsule to reconstruct the given input image:

with graph.as_default() as g:
    with tf.variable_scope('Masking'):
        
        #select the activity vector of given input image using the actual label y and mask out others
        masked_v = tf.multiply(tf.squeeze(digit_capsules), tf.reshape(y, (-1, 10, 1)))