Solved by ptrblck in post #2 
              
              
                The returned counts tensor c should have the shape [3] (not [1, 3]), which is also what I get. 
This is the expected shape, since the counts tensor will have the shape output.size(dim), if dim was specified. From the docs: 



counts (Tensor): (optional) if return_counts is True, there will be an ad…