(self, x)
| 133 | return x |
| 134 | |
| 135 | def forward(self, x): |
| 136 | # transform the input |
| 137 | x = self.stn(x) |
| 138 | |
| 139 | # Perform the usual forward pass |
| 140 | x = F.relu(F.max_pool2d(self.conv1(x), 2)) |
| 141 | x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)) |
| 142 | x = x.view(-1, 320) |
| 143 | x = F.relu(self.fc1(x)) |
| 144 | x = F.dropout(x, training=self.training) |
| 145 | x = self.fc2(x) |
| 146 | return F.log_softmax(x, dim=1) |
| 147 | |
| 148 | |
| 149 | model = Net().to(device) |