I have a code to generate 2D grid
B, C, H, W = x.size()
# mesh grid
xx = torch.arange(0, W).view(1,-1).repeat(H,1)
yy = torch.arange(0, H).view(-1,1).repeat(1,W)
xx = xx.view(1,1,H,W).repeat(B,1,1,1)
yy = yy.view(1,1,H,W).repeat(B,1,1,1)
grid = torch.cat((xx,yy),1).float()
How to obtain 3D grid by adding zz in the code? This is what I tried
B,C,D,H,W=input.size()
xx = torch.arange(0, W).view(1, 1,-1).repeat(D, H, 1)
yy = torch.arange(0, H).view(1, -1,1).repeat(D, 1, W)
zz = torch.arange(0, D).view(-1, 1,1).repeat(1,H,W)
xx = xx.view(1,1,D,H,W).repeat(B,1,1,1,1)
yy = yy.view(1,1,D,H,W).repeat(B,1,1,1,1)
zz = zz.view(1,1,D,H,W).repeat(B,1,1,1,1)
grid = torch.cat((xx,yy,zz),1).float().to('cuda')
But i feel it wrong in somewhere. Could you verify it?