import torch


def convert_2Djoints_to_gaussian_heatmaps_torch(joints2D,
                                                img_wh,
                                                std=4):
    """
    :param joints2D: (B, N, 2) tensor - batch of 2D joints.
    :param img_wh: int, dimensions of square heatmaps
    :param std: standard deviation of gaussian blobs
    :return heatmaps: (B, N, img_wh, img_wh) - batch of 2D joint heatmaps (channels first).
    """
    device = joints2D.device

    xx, yy = torch.meshgrid(torch.arange(img_wh, device=device),
                            torch.arange(img_wh, device=device))
    xx = xx[None, None, :, :].float()
    yy = yy[None, None, :, :].float()

    j2d_u = joints2D[:, :, 0, None, None]  # Horizontal coord (columns)
    j2d_v = joints2D[:, :, 1, None, None]  # Vertical coord (rows)
    heatmap = torch.exp(-(((xx - j2d_v) / std) ** 2) / 2 - (((yy - j2d_u) / std) ** 2) / 2)
    return heatmap