Extra changes, also not sure if works

model/stn.py

@@ -62,7 +62,7 @@ class STNet(nn.Module):
 
 
     # Spatial transformer network forward function
-    def stn(self, x):
+    def stn(self, x, t):
         xs = self.localization_net(x)
         xs = xs.view(-1, np.prod(self._loc_net_out_shape))  # -> (N, whatever)
         theta = self.fc_loc(xs)
@@ -71,12 +71,18 @@ class STNet(nn.Module):
         grid = F.affine_grid(theta, x.size())
         x = F.grid_sample(x, grid)
 
-        return x
+        # Do the same transformation to t sans training
+        with torch.no_grad():
+            t = t.view(t.size(0), 1, t.size(1), t.size(2))
+            t = F.grid_sample(t, grid)
+            t = t.squeeze(1)
+
+        return x, t
 
 
-    def forward(self, x):
+    def forward(self, x, t):
         # transform the input, do nothing else
-        return self.stn(x)
+        return self.stn(x, t)
 
 
 if __name__ == "__main__":