The easiest way to obtain a grid for use with the grid_sample() function is probably to use torch.nn.functional.affine_grid() 35.
You give it an affine matrix, and it returns a grid that you can then pass to grid_sample(). If you want an identity grid (no transformation), you can just pass it the identity affine matrix, which for 3D data is
aff = torch.FloatTensor([[[1, 0, 0, 0],[0, 1, 0, 0],[0, 0, 1, 0]]])

For example,
aff = aff.expand(B, 3, 4)  # expand to the number of batches you need
grid = torch.nn.functional.affine_grid(aff, size=(B,C,D,H,W))
torch.nn.functional.grid_sample(3d_data, grid)

So yeah, you could probably get the same effect using your code sample if you do it just right, but this should be much easier.
I hope this helps.