How to use affine_grid for batch tensor in pytorch?

Question

Following the official tutorial on affine_grid, this line (inside function stn()):

grid = F.affine_grid(theta, x.size())

gives me an error:

RuntimeError: Expected size for first two dimensions of batch2 tensor to be: [10, 3] but got: [4000, 3].

using the input as follow:

from torchinfo import summary
model = Net()
summary(model, input_size=(10, 1, 256, 256))

theta is of size (4000, 2, 3) and x is of size (10, 1, 256, 256), how do I properly manipulate theta for correct dimension to use with batch tensor?

EDIT: I honestly don't see any mistake here, and the dimension is actually according the the affine_grid doc, is there something changed in the function itself?