File size: 417 Bytes
57bdca5 |
1 2 3 4 5 |
We will actually write two: one that extracts the hidden features from a batch of images (like [BertModel]) and one that is suitable for image classification (like [BertForSequenceClassification]). As we mentioned before, we'll only write a loose wrapper of the model to keep it simple for this example. The only thing we need to do before writing this class is a map between the block types and actual block classes. |