harriskr14/garbage_classifier

This model has been pushed to the Hub using the PytorchModelHubMixin integration:

• Code: https://huggingface.co/harriskr14/garbage_classifier
• Paper: [More Information Needed]
• Docs: [More Information Needed]

How to use this model?

Write this following code on your file:

app.py

def build_model(num_classes=10):
    config = get_config()
    model_name = 'vit_base_patch16_224'
    model = timm.create_model(model_name, pretrained=True, num_classes=num_classes)
    return model

class MyModel(nn.Module, PyTorchModelHubMixin, repo_url="https://huggingface.co/harriskr14/garbage_classifier", pipeline_tag="text-classification", license="mit"):
    def __init__(self, num_classes=10):
        super(MyModel, self).__init__()
        self.model = build_model(num_classes=num_classes)

    def forward(self, x):
        return self.model(x)

model = MyModel()
model.from_pretrained("harriskr14/garbage_classifier")