upload source of model file main.py

main.py

@@ -0,0 +1,58 @@
+import torch
+from sklearn.datasets import load_iris
+from torch.utils.data import DataLoader, Dataset
+import torch.nn as nn
+import torch.nn.functional as fnn
+import torch.optim as optim
+
+class IrisDataset(Dataset):
+    def __init__(self):
+        data = load_iris()
+        self.data = torch.tensor(data.data, dtype=torch.float32)
+        self.targets = torch.tensor(data.target, dtype=torch.long)
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        return self.data[idx], self.targets[idx]
+
+dataset = IrisDataset()
+dataloader = DataLoader(dataset, batch_size=16, shuffle=True)
+
+class INet(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc1 = nn.Linear(4, 64)
+        self.fc2 = nn.Linear(64, 32)
+        self.fc3 = nn.Linear(32, 16)
+        self.fc4 = nn.Linear(16, 8)
+        self.fc5 = nn.Linear(8, 4)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = fnn.relu(x)
+        x = self.fc2(x)
+        x = fnn.relu(x)
+        x = self.fc3(x)
+        x = fnn.relu(x)
+        x = self.fc4(x)
+        x = fnn.relu(x)
+        x = self.fc5(x)
+        return x
+
+model = INet()
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model.parameters(), lr=0.01)
+
+for epoch in range(30):
+    for X, y in dataloader:
+        preds = model(X)
+        loss = criterion(preds, y)
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+    print(f"Epoch {epoch}, Loss: {loss.item()}")
+
+torch.save(model.state_dict(), "/models/inet.pth")