Update app.py

app.py

@@ -1,40 +1,78 @@
+# import torch
+
+# # Utility to log tensor info
+# def log_tensor(name, x):
+#     print(f"--- {name} ---")
+#     print(f"shape: {x.shape}, dtype: {x.dtype}, device: {x.device}")
+#     print(f"min: {x.min().item():.6f}, max: {x.max().item():.6f}, mean: {x.mean().item():.6f}, sum: {x.sum().item():.6f}")
+#     print(f"full tensor:\n{x}\n")
+
+# # Simple function
+# def g(x, y):
+#     z = x + y
+#     return z
+
+# # Compiled function
+# @torch.compile(backend="eager")
+# def f(x):
+#     x = torch.sin(x)
+#     x = g(x, x)
+#     return x
+
+# # --- Initialization / run once ---
+# def init_and_run_once():
+#     # Example input
+#     x = torch.ones(3, 3, dtype=torch.float32)
+#     print("=== INITIAL INPUT ===")
+#     log_tensor("original input x", x)
+
+#     # Run compiled function once
+#     out = f(x)
+
+#     print("=== OUTPUT AFTER FIRST RUN ===")
+#     log_tensor("final output", out)
+
+#     # Return output if needed
+#     return out
+
+# # Run once at import / init
+# if __name__ == "__main__":
+#     init_and_run_once()
+
+import os
 import torch
+from torch import nn
+from torch.utils.data import DataLoader
+from torchvision import datasets, transforms
+device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
+print(f"Using {device} device")
+
+class NeuralNetwork(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.flatten = nn.Flatten()
+        self.linear_relu_stack = nn.Sequential(
+            nn.Linear(28*28, 512),
+            nn.ReLU(),
+            nn.Linear(512, 512),
+            nn.ReLU(),
+            nn.Linear(512, 10),
+        )
+
+    def forward(self, x):
+        x = self.flatten(x)
+        logits = self.linear_relu_stack(x)
+        return logits
+
+
+#We create an instance of NeuralNetwork, and move it to the device, and print its structure.
+model = NeuralNetwork().to(device)
+print(model)
+
+#To use the model, we pass it the input data. This executes the model’s forward, along with some background operations. 
 
-# Utility to log tensor info
-def log_tensor(name, x):
-    print(f"--- {name} ---")
-    print(f"shape: {x.shape}, dtype: {x.dtype}, device: {x.device}")
-    print(f"min: {x.min().item():.6f}, max: {x.max().item():.6f}, mean: {x.mean().item():.6f}, sum: {x.sum().item():.6f}")
-    print(f"full tensor:\n{x}\n")
-
-# Simple function
-def g(x, y):
-    z = x + y
-    return z
-
-# Compiled function
-@torch.compile(backend="eager")
-def f(x):
-    x = torch.sin(x)
-    x = g(x, x)
-    return x
-
-# --- Initialization / run once ---
-def init_and_run_once():
-    # Example input
-    x = torch.ones(3, 3, dtype=torch.float32)
-    print("=== INITIAL INPUT ===")
-    log_tensor("original input x", x)
-
-    # Run compiled function once
-    out = f(x)
-
-    print("=== OUTPUT AFTER FIRST RUN ===")
-    log_tensor("final output", out)
-
-    # Return output if needed
-    return out
-
-# Run once at import / init
-if __name__ == "__main__":
-    init_and_run_once()
+X = torch.rand(1, 28, 28, device=device)
+logits = model(X)
+pred_probab = nn.Softmax(dim=1)(logits)
+y_pred = pred_probab.argmax(1)
+print(f"Predicted class: {y_pred}")