Premium problem36. Sum Over Last Two Axes

Medium Locked

Given a 4-D tensor A of shape (a, b, c, d), return the sum over its last two dimensions, producing a 2-D tensor of shape (a, b).

Input

tensor([[[[  0,  ...,   3],
          ...,
          [  8,  ...,  11]],

         ...,

         [[ 36,  ...,  39],
          ...,
          [ 44,  ...,  47]]],


        ...,


        [[[ 96,  ...,  99],
          ...,
          [104,  ..., 107]],

         ...,

         [[132,  ..., 135],
          ...,
          [140,  ..., 143]]]])

Output

tensor([[  66,  210,  354,  498],
        [ 642,  786,  930, 1074],
        [1218, 1362, 1506, 1650]])

Premium problem

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

Implement solve(...)