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?