import torch


def _get_size(tensor):
    """Returns the size of a tensor, but treats a scalar as a 1-element vector"""
    if tensor.dim() == 0:
        return torch.Size([1])
    else:
        return tensor.size()