heheyas
init
cfb7702
raw
history blame
745 Bytes
import torch
import torch.nn as nn
from utils.misc import get_rank
class BaseModel(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.rank = get_rank()
self.setup()
if self.config.get('weights', None):
self.load_state_dict(torch.load(self.config.weights))
def setup(self):
raise NotImplementedError
def update_step(self, epoch, global_step):
pass
def train(self, mode=True):
return super().train(mode=mode)
def eval(self):
return super().eval()
def regularizations(self, out):
return {}
@torch.no_grad()
def export(self, export_config):
return {}