import json,torch,yaml
from pathlib import Path
from torch import nn
def cfg(r):return yaml.safe_load((Path(r)/'conf/config.yaml').read_text())
class WB2Baseline(nn.Module):
 def __init__(self,variables=8,hidden=16):super().__init__();self.n=nn.Sequential(nn.Conv2d(variables,hidden,3,padding=1),nn.ReLU(),nn.Conv2d(hidden,variables,3,padding=1));self.model_config={'variables':variables,'hidden':hidden}
 def forward(self,x):return x+self.n(x)
def write(p,o):p=Path(p);p.parent.mkdir(parents=True,exist_ok=True);p.write_text(json.dumps(o,indent=2)+'\n')