Premium problem27. Multiply (5,5,3) by (5,5)

Medium Locked

Multiply (5, 5, 3) tensor A by (5, 5) tensor B, broadcasting B over the last axis. Return the result.

Input

A =
tensor([[[1.,  ..., 1.],
         ...,
         [1.,  ..., 1.]],

        ...,

        [[1.,  ..., 1.],
         ...,
         [1.,  ..., 1.]]])
B =
tensor([[2., 2., 2., 2., 2.],
        [2., 2., 2., 2., 2.],
        [2., 2., 2., 2., 2.],
        [2., 2., 2., 2., 2.],
        [2., 2., 2., 2., 2.]])

Output

tensor([[[2.,  ..., 2.],
         ...,
         [2.,  ..., 2.]],

        ...,

        [[2.,  ..., 2.],
         ...,
         [2.,  ..., 2.]]])

Premium problem

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

Implement solve(...)