Upload model.py with huggingface_hub

model.py

@@ -0,0 +1,18 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+class Net(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv2d(16, 8, kernel_size=3, padding=1)
+        self.fc1 = nn.Linear(8 * 64 * 64, 32)  # for 256x256 input
+        self.fc2 = nn.Linear(32, 4)  # 4 classes
+
+    def forward(self, x):
+        out = F.max_pool2d(torch.tanh(self.conv1(x)), 2)  # 256 -> 128
+        out = F.max_pool2d(torch.tanh(self.conv2(out)), 2)  # 128 -> 64        
+        out = out.view(-1, 8 * 64 * 64)
+        out = torch.tanh(self.fc1(out))
+        out = self.fc2(out)
+        return out