epochs test

app.py

@@ -1,14 +1,52 @@
 import gradio as gr
 import spaces
 import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
 
-zero = torch.Tensor([0]).cuda()
-print(zero.device) # <-- 'cpu' 🤔
+class SimpleNet(nn.Module):
+    def __init__(self):
+        super(SimpleNet, self).__init__()
+        self.fc = nn.Linear(784, 10)  # Simple model for MNIST
+
+    def forward(self, x):
+        x = x.view(-1, 784)  # Flatten the image
+        x = self.fc(x)
+        return x
 
 @spaces.GPU
-def greet(n):
-    print(zero.device) # <-- 'cuda:0' 🤗
-    return f"Hello {zero + n} Tensor"
+def train_model(epochs):
+    # Load MNIST dataset
+    transform = transforms.Compose([transforms.ToTensor()])
+    train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
+    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
+
+    # Model, loss, and optimizer
+    model = SimpleNet().cuda()
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
+
+    # Training loop
+    for epoch in range(epochs):
+        model.train()
+        running_loss = 0.0
+        for data, target in train_loader:
+            data, target = data.cuda(), target.cuda()
+            optimizer.zero_grad()
+            output = model(data)
+            loss = criterion(output, target)
+            loss.backward()
+            optimizer.step()
+            running_loss += loss.item()
+
+        print(f"Epoch {epoch + 1}, Average Loss: {running_loss / len(train_loader)}")
+    
+    # Save the model checkpoint
+    torch.save(model.state_dict(), "simple_net.pth")
+    return "Training completed and model saved."
 
-demo = gr.Interface(fn=greet, inputs=gr.Number(), outputs=gr.Text())
-demo.launch()
+# Define the Gradio interface
+demo = gr.Interface(fn=train_model, inputs=gr.Slider(1, 5, step=1, default=1, label="Number of Epochs"), outputs="text")
+demo.launch()