afaf4383126b25be6ed427c0f3626ae09ae35b45,torchsample/utils.py,,th_meshgrid,#,61

Before Change

def th_meshgrid(*args):
    pools = (torch.range(0,i-1) for i in args)
    result = [[]]
    for pool in pools:
        result = [x+[y] for x in result for y in pool]
    return torch.Tensor(result).long()

After Change

    dtype = torch.LongTensor
    pools = []
    for i in args:
        if isinstance(i, int):
            pools.append(torch.range(0, i-1))
        else:
            if type(i) != torch.LongTensor:
                dtype = i.type()
            pools.append(i)
    result = [[]]
    for pool in pools:
        result = [x+[y] for x in result for y in pool]
    return torch.Tensor(result).type(dtype)

Instances