Generating histogram feature of 2D tensor from 3D Tensor feature set

1

I have a 3D tensor of dimensions (3,4 7) where each element in 2-dim(4) has 7 attributes.
What I want is to take the 4th attribute of all 4 elements and to calculate the histogram having 3
hist values and store those values only. And ending up with a 2D tensor of shape (3,4).
I have a small toy example for the task that I am working on. My solution ends up with a Tensor which has shape (1,3). Any hint or guidance will be appreciated

import torch
torch.manual_seed(1)
feature = torch.randint(1, 50, (3, 4,7))
feature.type(torch.FloatTensor)
attrbute_val = feature[:,:,3:4]
print(attrbute_val.shape)
print(attrbute_val)
histogram_feature = torch.histc(torch.tensor(attrbute_val,dtype=torch.float32), bins=3, min=1, max=50)
print("histogram_feature",histogram_feature)
2

This is want i have come up with a naive solution

attrbute_val = feature[:,:,3:4]
print(attrbute_val.shape)
print(attrbute_val[:,:,0])

final_feature = np.zeros((3,3))
shape = attrbute_val[:,:,0].shape
for row in range(shape[0]):
  print("element wise features",attrbute_val[:,:,0][row])
  hist,_= np.histogram(attrbute_val[:,:,0][row], bins=3)
  print(row,",hist values",hist)
  final_feature[row,:] = hist

print("final feature shape", final_feature.shape)
print("final feature",final_feature)

I wish PyTorch has any way to apply customer function on dimensions