import torch


def pose_square(pose):
    """Converts pose matrix of size 3x4 to a square matrix of size 4x4"""

    pose_sh = pose.shape
    if pose_sh[2] == 3:
        pose_row = torch.tensor([0., 0., 0., 1.])
        if pose.is_cuda:
            pose_row = pose_row.to(pose.device)
        pose_row = pose_row.repeat(pose_sh[0], pose_sh[1], 1, 1)
        pose = torch.cat((pose, pose_row), 2)

    return pose