Swap Height and Width

easy (<10 mins) einops tensor manipulation image processing

this year by E

Description

For a batch of images or feature maps, swap the height and width dimensions. The input shape is (B, C, H, W) and the output shape should be (B, C, W, H).

Starter Code

import torch
from einops import rearrange

def swap_height_width(tensor):
    # Your einops code here
    pass

Verification

For an input of shape (10, 3, 224, 256), the output should have the shape (10, 3, 256, 224).