(self, x)
| 218 | return nn.Sequential(*layers) |
| 219 | |
| 220 | def _forward_impl(self, x): |
| 221 | # See note [TorchScript super()] |
| 222 | x = self.conv1(x) |
| 223 | x = self.bn1(x) |
| 224 | x = self.relu(x) |
| 225 | x = self.maxpool(x) |
| 226 | |
| 227 | x = self.layer1(x) |
| 228 | x = self.layer2(x) |
| 229 | x = self.layer3(x) |
| 230 | x = self.layer4(x) |
| 231 | |
| 232 | x = self.avgpool(x) |
| 233 | x = torch.flatten(x, 1) |
| 234 | out = self.fc(x) |
| 235 | if self.is_remix: |
| 236 | rot_output = self.rot_classifier(x) |
| 237 | return out, rot_output |
| 238 | else: |
| 239 | return out |
| 240 | |
| 241 | def forward(self, x): |
| 242 | return self._forward_impl(x) |