How to properly use affine_grid for batch tensor?

1

I was following this official (I think?) tutorial on Spatial Transformer Network, but the affine_grid seems to be error-nous. I get error saying

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

on the line

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

by simply running:

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

What’s the catch here? Thanks in advance.