How to resize an affine_grid

Hi guys,
I have a simple question: normally when I want to resize a 4D tensor, I apply the resize method provided by the torchvision.transforms library, and everything works correctly because I pass a 4D tensor with this form: [num_batches, num_channels, Width, Height].

I wanted to try the same function with the output of the affine_grid function and I found that the output has this form: [num_channels, Width, Height, 2]

Could someone give me an idea or example on how to resize the output of F.affine_grid?

Thank you in advance

just implement the function nn.permute then resize the image using resize then reuse the method permute to set the output of the affine_grid as the original form.
transform2 = T.Resize((28,28))
grid = grid.permute(0, 3, 1, 2)
grid = transform2(grid)
grid = grid.permute(0, 2, 3, 1)