@fmassa
The relevant code is:

batch=torch.from_numpy(batch) #batch is the matrix of a image ranged from 0 to 1 with batch size 3, size(3,1200,1600,3)
initial=np.array([[2,0,0],[0,2,0]])
initial=initial.astype('float32')
initial=initial.flatten()
W=torch.zeros(1200*1600*3,6)
b=torch.Tensor(initial)
h=torch.from_numpy(torch.zeros(3,1200*1600*3).numpy() @ W.numpy() +b.numpy())
h=h.view(-1,2,3).type(torch.Tensor) #size(3,2,3)
x_t=torch.mm(torch.ones(600,1),torch.unsqueeze(torch.linspace(-1.0,1.0,800),1).transpose(1,0))  y_t=torch.mm(torch.unsqueeze(torch.linspace(-1.0,1.0,600),1),torch.ones(1,800))
x_t_flat=x_t.view(1,-1)
y_t_flat=y_t.view(1,-1)
ones=torch.ones(x_t_flat.size())
grid=torch.cat([x_t_flat,y_t_flat,ones],0)
grid=grid.view(-1).repeat(3,1).view(3,3,-1) #size(3,3,480000)

First,I use this:

T=torch.from_numpy(h.numpy() @ grid.numpy())

and then another way:

T_another=torch.bmm(h,grid)

the subsequent code is omitted.
and I have the last output of the image matrix that can’t imshow() normally with the second way but the first is OK. When I look back on the code, I just found the T is unequal to T_another.