Model Card for ombayus_iris_model

import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
  def __init__(self,in_features=4,h1=8,h2=9,out_feauteres=3):
    super().__init__() # instantiate our nn.Module
    self.fc1 = nn.Linear(in_features,h1)
    self.fc2 = nn.Linear(h1,h2)
    self.out = nn.Linear(h2,out_feauteres)

  def forward(self,x):
    x = F.relu(self.fc1(x))
    x = F.relu(self.fc2(x))
    x = self.out(x)
    return x

def number_to_follower(x):
  if x == 0:
    return'Setosa'
  elif x == 1:
    return 'Versicolor'
  elif x == 2:
    return 'Virginica'

model = Model()

model.load_state_dict(torch.load('ombayus_iris_model.pt'))

new_iris = torch.tensor([5.9,3.0,5.1,1.8])

with torch.no_grad():
  print(number_to_follower(model(new_iris).argmax().item()))