text: # WARNING: you are on the master branch, please refer to the examples on the branch that matches your `cortex version`
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from torch.autograd import Variable
  from sklearn.datasets import load_iris
  from sklearn.model_selection import train_test_split
  from sklearn.metrics import accuracy_score
  
  
  class IrisNet(nn.Module):
      def __init__(self):
          super(IrisNet, self).__init__()
          self.fc1 = nn.Linear(4, 100)
          self.fc2 = nn.Linear(100, 100)
          self.fc3 = nn.Linear(100, 3)
          self.softmax = nn.Softmax(dim=1)
  
      def forward(self, X):
          X = F.relu(self.fc1(X))
          X = self.fc2(X)
          X = self.fc3(X)
          X = self.softmax(X)
          return X
  
  
  if __name__ == "__main__":
      iris = load_iris()
      X, y = iris.data, iris.target
      X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.8, random_state=42)
  
      train_X = Variable(torch.Tensor(X_train).float())
      test_X = Variable(torch.Tensor(X_test).float())
      train_y = Variable(torch.Tensor(y_train).long())
      test_y = Variable(torch.Tensor(y_test).long())
  
      model = IrisNet()
  
      criterion = nn.CrossEntropyLoss()
  
      optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
  
      for epoch in range(1000):
          optimizer.zero_grad()
          out = model(train_X)
          loss = criterion(out, train_y)
          loss.backward()
          optimizer.step()
  
          if epoch % 100 == 0:
              print("number of epoch {} loss {}".format(epoch, loss))
  
      predict_out = model(test_X)
      _, predict_y = torch.max(predict_out, 1)
  
      print("prediction accuracy {}".format(accuracy_score(test_y.data, predict_y.data)))
  
      torch.save(model.state_dict(), "weights.pth")
  