def shrink_batch(batch):
    """
    Exclude certain keys from the batch that take up a lot of memory
    """

    bad_keys = ['nbr_list', 'kj_idx', 'ji_idx',
                'angle_list']
    new_batch = {key: val for key, val in batch.items()
                 if key not in bad_keys}

    return new_batch