import torch


def get_closest_color(a, b):
    """
    a: h1 x w1 x 3 pytorch tensor
    b: h2 x w2 x 3 pytorch tensor
    returns: h1 x w1 x 3 pytorch tensor containing the nearest color in b to the corresponding pixels in a"""
    a_flat = a.contiguous().view(1, -1, 3)
    b_flat = b.contiguous().view(-1, 1, 3)
    a_b = torch.norm(a_flat - b_flat, dim=2, keepdim=True)
    index = torch.argmin(a_b, dim=0)
    closest_color = b_flat[index]
    return closest_color.contiguous().view(a.shape)