For anyone else stuck with similar problem, the reason input of variable size doesn’t work is because torch.stack will fail for inputs of varying size. Their is also a github issue in this regard. One way to work with such input is to write custom collate function. For example, see below.
def collate_fn_custom( batch):
    images = list()
    boxes = list()
    objectLabels = list()
    distances = list()

    for b in batch:
        images.append(b[0])
        objectLabels.append(b[1])
        boxes.append(b[2])
        
        distances.append(b[3])

    images = torch.stack(images, dim=0)

    return images, objectLabels, boxes, distances

Related issue and source :https://github.com/pytorch/pytorch/issues/1512 2