k = torch.arange(a.size(0))[:, None] + torch.arange(b.size(0))[None, :]
kk = k.new()
kk.set_(k.storage(), 0, (k.size(0), k.size(0) + k.size(1) - 1), (k.size(1) - 1, 1))
k = torch.arange(a.size(0))[:, None] + torch.arange(b.size(0))[None, :]
kk = k.new()
kk.set_(k.storage(), 0, (k.size(0), k.size(0) + k.size(1) - 1), (k.size(1) - 1, 1))