Args: inputs: input should have spatially N dimensions ``(Batch, in_channels, dim_0[, dim_1, ..., dim_N])``, N is defined by `dimensions`. Returns: a torch Tensor of classification prediction in shape ``(Batch, num_classes)``.
(self, inputs: torch.Tensor)
| 423 | block.set_swish(memory_efficient) |
| 424 | |
| 425 | def forward(self, inputs: torch.Tensor): |
| 426 | """ |
| 427 | Args: |
| 428 | inputs: input should have spatially N dimensions |
| 429 | ``(Batch, in_channels, dim_0[, dim_1, ..., dim_N])``, N is defined by `dimensions`. |
| 430 | |
| 431 | Returns: |
| 432 | a torch Tensor of classification prediction in shape ``(Batch, num_classes)``. |
| 433 | """ |
| 434 | # Stem |
| 435 | x = self._conv_stem(self._conv_stem_padding(inputs)) |
| 436 | x = self._swish(self._bn0(x)) |
| 437 | # Blocks |
| 438 | x = self._blocks(x) |
| 439 | # Head |
| 440 | x = self._conv_head(self._conv_head_padding(x)) |
| 441 | x = self._swish(self._bn1(x)) |
| 442 | |
| 443 | # Pooling and final linear layer |
| 444 | x = self._avg_pooling(x) |
| 445 | |
| 446 | x = x.flatten(start_dim=1) |
| 447 | x = self._dropout(x) |
| 448 | x = self._fc(x) |
| 449 | return x |
| 450 | |
| 451 | def _initialize_weights(self) -> None: |
| 452 | """ |