Here is the rest of the code, thank you.
class MLP(nn.Module):
    def __init__(self, sz):
        super(MLP, self).__init__()
        self.layer1 = nn.Linear(sz, 100)
        self.layer2 = nn.Linear(100,2)

    def forward(self, x):
        x = relu(self.layer1(x))
        x = self.layer2(x)
        return x

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.net = MLP(28)
        self.opt = torch.optim.SGD(self.parameters(), lr=0.1)
        self.bce = torch.nn.CrossEntropyLoss()

    def observe(self, x, y):
        self.train()
        self.zero_grad()
        y_ = self.net(x)
        loss = self.bce(y_, y)
        loss.backward()
        self.opt.step()

        _, idx = torch.max(y_.data.cpu(), 1 , keepdim=False)
        acc = (idx==y.data.cpu()).float()
        return loss, acc