File size: 496 Bytes
fe3b346 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | 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
|