def METHOD_NAME(self, tensor):
    size = tensor.size()
    if len(size) >= 2:
        batch_size = size[0] * size[1]
        expanded_size = (
            (batch_size,) + size[2:] if len(size) > 2
            else (batch_size,)
        )
        tensor = tensor.view(expanded_size)
    return tensor