text_plain_output_1.png
from torch import tensor, nn

x_train, y_train, x_valid, y_valid = get_data()
n, m = x_train.shape
c = y_train.max() + 1
nh = 50
(n, m, c, nh)

class Model(nn.Module):
    def __init__(self, ni, nh, no):
        super().__init__()
        self.layers = [nn.Linear(ni, nh), nn.ReLU(), nn.Linear(nh, no)]
    def __call__(self, x):
        for l in self.layers:
            x = l(x)
        return x
model = Model(m, nh, 10)
pred = model(x_train)
pred.shape