vancauwe's picture
feat: template for model on hugging face
fe3b346 verified
raw
history blame contribute delete
496 Bytes
from typing import Dict, List, Optional, Tuple
import timm
import torch
from pytorch_lightning import LightningModule
class TemplateClassifier(LightningModule):
def __init__(self, config: dict):
super().__init__()
# NN architecture
self.backbone = timm.create_model(
#SPECIFY HERE YOUR MODEL
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# WRITE YOU CODE HERE
predictions=None
return predictions