def pack(self, grads):
        """
        Args:
            grads (list): list of gradient tensors

        Returns:
            packed list of gradient tensors to be aggregated.
        """
        for i, g in enumerate(grads):
            assert g.shape == self._shapes[i]

        with cached_name_scope("GradientPacker", top_level=False):
            concat_grads = tf.concat([tf.reshape(g, [-1]) for g in grads], 0, name='concatenated_grads')
            # concat_grads = tf.cast(concat_grads, tf.float16)
            grad_packs = tf.split(concat_grads, self._split_sizes)
            return grad_packs

        Args:
            grads (list): list of gradient tensors

        Returns:
            packed list of gradient tensors to be aggregated.
        
tensorpack/graph_builder/utils.py