Rotate a Tensor in PyTorch

rot.py

import numpy as np

import torch
import torch.nn.functional as F

def rotz(theta):
    return np.array([
        [np.cos(theta), -np.sin(theta), 0],
        [np.sin(theta), np.cos(theta), 0],
        [0, 0, 1],
    ], dtype="float64")

device = torch.device("cuda")
x = torch.rand(5, 3, 200, 200).to(device)
angle = np.deg2rad(7)
rotm = rotz(angle)[:2, :].reshape(2, 3, 1)
rotm = torch.FloatTensor(rotm).permute(2, 0, 1).repeat(x.shape[0], 1, 1)
affine_grid = F.affine_grid(rotm, x.size()).to(device)
with torch.no_grad():
    x_r = F.grid_sample(x, affine_grid)