import tensorflow as tf

                        W_d = W_d[input_idx]
                        W_p = self._make_var('W_p', (ni, 1, 1, ch, ch))
                        W_p = W_p[input_idx]
                        batch_norm_offset = self._make_var('batch_norm_offset', (ni, ch), init_constant=0)
                        batch_norm_offset = batch_norm_offset[input_idx]
                        batch_norm_scale = self._make_var('batch_norm_scale', (ni, ch), init_constant=1)
                        batch_norm_scale = batch_norm_scale[input_idx]

                    X = self._do_separable_conv(X, w, h, ch, filter_size=filter_size, stride=stack_stride,
                                                W_d=W_d, W_p=W_p, no_batch_norm=True)
                    X = self._add_batch_norm(X, ch, offset=batch_norm_offset, scale=batch_norm_scale,
                                            no_moving_average=is_dynamic, is_train=is_train)

        X = tf.reshape(X, (-1, w // stride, h // stride, ch)) # Sanity shape check
        return X

    ####################################
    # Utils
    ####################################

    def _do_cutout(self, image, im_width, im_height, cutout_size):
        mask = tf.ones([cutout_size, cutout_size], dtype=tf.int32)
        start_x = tf.random.uniform(shape=(1,), minval=0, maxval=im_width, dtype=tf.int32)
        start_y = tf.random.uniform(shape=(1,), minval=0, maxval=im_height, dtype=tf.int32)
        mask = tf.pad(mask, [[cutout_size + start_y[0], im_height - start_y[0]],
                            [cutout_size + start_x[0], im_width - start_x[0]]])
