Premium problem49. Stack Tensors Along a New Axis

Easy Locked

Stack a and b along a new dimension inserted at position dim. This adds an extra dimension to the output, rather than joining the tensors along an existing one.

Input

a = tensor([1, 2, 3, 4])
b = tensor([5, 6, 7, 8])
dim = 1

Output

tensor([[1, 5],
        [2, 6],
        [3, 7],
        [4, 8]])

Premium problem

This one's part of Premium. Unlock the full PyTorch track plus every other premium problem on the site.

Implement solve(...)