Hi,
If you are in the mood for a weekend hack, here is a not entirely serious and apparently not very stable (it seemed to be crashing for me fairly often) solution without a for loop.
I’m not sure how you do with more than very few features, but here you go.
It uses the fact that sparse matrix indices may contain the same coordinate multiple times, the matrix entry is then the sum of all values at the coordinate.
import torch
from matplotlib import pyplot
%matplotlib inline

data = torch.randn(2,1000,2) # batch x draws x (x,y)
data[1,:,1] += data[1,:,0]
d_scaled = (data*2.5).long().view(-1,2)
d_scaled -= d_scaled.min()
d_idx0 = (torch.arange(0,data.size(0)).long().view(-1,1)*torch.LongTensor(1,data.size(1)).fill_(1)).view(-1,1)
d_idx = torch.cat([d_idx0, d_scaled], dim=1)
d_ones = torch.FloatTensor(d_idx0.size(0)).fill_(1.0)
st = torch.sparse.FloatTensor(d_idx.t(),d_ones,torch.Size((2,20,20)))
hist = st.to_dense()
pyplot.subplot(1,2,1)
pyplot.contour(hist[0].numpy(), extent=(-4,4,-4,4))
pyplot.subplot(1,2,2)
pyplot.contour(hist[1].numpy(), extent=(-4,4,-4,4))

Best regards
Thomas