Now I have found a strange thing,
I have a tensor theta with the size (3,2,3)

theta


2  0  0
0  2  0
[torch.FloatTensor of size 2x3]
2  0  0
0  2  0
[torch.FloatTensor of size 2x3]
2  0  0
0  2  0
[torch.FloatTensor of size 2x3]

and a tensor grid with the size (3,3,480000)

grid


-1.0000 -0.9975 -0.9950  …   0.9950  0.9975  1.0000
-1.0000 -1.0000 -1.0000  …   1.0000  1.0000  1.0000
1.0000  1.0000  1.0000  …   1.0000  1.0000  1.0000
[torch.FloatTensor of size 3x480000]
-1.0000 -0.9975 -0.9950  …   0.9950  0.9975  1.0000
-1.0000 -1.0000 -1.0000  …   1.0000  1.0000  1.0000
1.0000  1.0000  1.0000  …   1.0000  1.0000  1.0000
[torch.FloatTensor of size 3x480000]
-1.0000 -0.9975 -0.9950  …   0.9950  0.9975  1.0000
-1.0000 -1.0000 -1.0000  …   1.0000  1.0000  1.0000
1.0000  1.0000  1.0000  …   1.0000  1.0000  1.0000
[torch.FloatTensor of size 3x480000]

then I use torch.bmm:

torch.bmm(theta,grid)


(  0   ,.,.) =
-2.0000 -1.9950 -1.9900  …   0.0000  0.0000  0.0000
-2.0000 -2.0000 -2.0000  …   0.0000  0.0000  0.0000


(  1   ,.,.) =
-2.0000 -1.9950 -1.9900  …   0.0000  0.0000  0.0000
-2.0000 -2.0000 -2.0000  …   0.0000  0.0000  0.0000


(  2   ,.,.) =
-2.0000 -1.9950 -1.9900  …   0.0000  0.0000  0.0000
-2.0000 -2.0000 -2.0000  …   0.0000  0.0000  0.0000
[torch.FloatTensor of size 3x2x480000]

and obviously the output is wrong, the right answer is:

theta.numpy() @ grid.numpy()


(  0   ,.,.) =
-2.0000 -1.9950 -1.9900  …   1.9900  1.9950  2.0000
-2.0000 -2.0000 -2.0000  …   2.0000  2.0000  2.0000


(  1   ,.,.) =
-2.0000 -1.9950 -1.9900  …   1.9900  1.9950  2.0000
-2.0000 -2.0000 -2.0000  …   2.0000  2.0000  2.0000


(  2   ,.,.) =
-2.0000 -1.9950 -1.9900  …   1.9900  1.9950  2.0000
-2.0000 -2.0000 -2.0000  …   2.0000  2.0000  2.0000
[torch.FloatTensor of size 3x2x480000]

so I wonder how it is happening?