Embed Co-ordinate Tensor in Array Tensor

Suppose I have a PyTorch tensor that represents a list of 1000 coordinates, where indices 0-2 are independent x-y-z co-ordinates and index 3 is a value associated with the coordinate.

I want to produce a ‘graph’ of those co-ordinates: an I x J x K grid, with the value from the coordinates at each index.

For example:

channels = 2
n_coords = 1000

len_coords = 4
low = 0
some_limit = 5

coords = torch.zeros(batch_size, channels, n_coords, len_coords).random_(low, high) # a set of coordinates
image = torch.zeros(batch_size, channels, high, high, high) #The 'space' I want to embed the coo-ordinates in

I have a solution involving loops that I have developed that produces the desired result:

coords_shape = torch.Tensor([*data.shape][0:2]).int()
axis_shape = torch.Tensor([high,high,high]).int()    
image_shape = torch.cat((coords_shape, axis_shape), axis=0)
image = torch.zeros(*image_shape)

for b in range(coords.shape[0]):
    for c in range(coords.shape[1]):
        for v in range(coords.shape[2]):
                elem = coords[b, c, v]
                    image[batch_elem, channel, int(elem[0]), int(elem[1]), int(elem[2])]=elem[3]

I know this is suboptimal because of the loops.

Is there a smarter way of implementing this?