<reponame>NullConvergence/repulsive_proto<gh_stars>0
import torch
import torch.nn as nn


class PreprocessingModel(nn.Module):
    def __init__(self, net, preproc=None):
        super().__init__()
        self.net = net
        self.preproc_args = preproc

    def preprocess(self, x):
        if self.preproc_args is None:
            return x
        else:
            return (x-self.preproc_args['mean']) / self.preproc_args['std']

    def forward(self, x):
        y = self.preprocess(x)
        return self.net(y)

    def train(self):
        self.net.train()

    def eval(self):
        self.net.eval()

import torch
import torch.nn as nn


class PreprocessingModel(nn.Module):
    def __init__(self, net, preproc=None):
        super().__init__()
        self.net = net
        self.preproc_args = preproc

    def preprocess(self, x):
        if self.preproc_args is None:
            return x
        else:
            return (x-self.preproc_args['mean']) / self.preproc_args['std']

    def forward(self, x):
        y = self.preprocess(x)
        return self.net(y)

    def train(self):
        self.net.train()

    def eval(self):
        self.net.eval()