Class linMod

models/linear.py:8–18

class linMod(nn.Module):
    def __init__(self, dim, n_class):
        super(linMod, self).__init__()
        self.dim = dim
        self.lm = nn.Linear(int(np.prod(dim)), n_class)
    def forward(self, x):
        x = x.view(-1, int(np.prod(self.dim)))
        out = self.lm(x)
        return out, x
    def get_embedding_dim(self):
        return int(np.prod(self.dim))

# mlp model class
class mlpMod(nn.Module):