Search code examples
pythondictionarytorch

Concatenate all 2 dimensional values in a dictionary. (Output is Torch tensor)


I want to concatenate all 2 dimensional values in a dictionary.

The number of rows of these values is always the same.

D = {'a': [[0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0]],
     'b': [[1, 1],
           [1, 1],
           [1, 1]],
     'c': [[2, 2, 2, 2],
           [2, 2, 2, 2],
           [2, 2, 2, 2]]
}

And the output must be form of a torch tensor.

tensor([[0, 0, 0, 0, 0, 1, 1, 2, 2, 2, 2],
        [0, 0, 0, 0, 0, 1, 1, 2, 2, 2, 2],
        [0, 0, 0, 0, 0, 1, 1, 2, 2, 2, 2]])

Any help would be appreciated!!


Solution

  • import torch
    print(torch.cat(tuple([torch.tensor(D[name]) for name in D.keys()]), dim=1))
    

    Output:

    tensor([[0, 0, 0, 0, 0, 1, 1, 2, 2, 2, 2],
            [0, 0, 0, 0, 0, 1, 1, 2, 2, 2, 2],
            [0, 0, 0, 0, 0, 1, 1, 2, 2, 2, 2]])