import torch


def combine_obj_pixels(obj_pix, obj_dim):
    """Combine obj-split pixels into a single image.
    Args:
        obj_pix: B, ..., Nobj, ..., C, H, W
        obj_dim: The dimension to reduce over -- which corresponds to objs
    Returns
        B, ..., ..., C, H, W
    """
    if obj_pix is None:
        return None
    return torch.max(obj_pix, dim=obj_dim)[0]