Hi, I am trying to implement stn. The affine_grid in torch.nn supports global transformation (for the whole image), i.e theta is Nx2x3 but I am looking for an implementation where each pixel has a 2x3 affine transformation. So theta in my case would be NxHxWx2x3.

Oh, so you want to transform each pixel independently! Sorry, I do not have more experience beyond this, but also I cannot imagine what would happen to an image if each pixel is transformed separately!